@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e4d61d0d22

[CPU] Add AMX-only high-performance MLA backend for DeepSeek V2/V3/R1 (#52616)

Signed-off-by: jiang1.li <jiang1.li@intel.com>

Li, Jiang committed Aug 19, 2026 at 10:28 UTC e4d61d0d222ad21437a1507bb34ca7e916aac880
20 files changed +4156 -6
.buildkite/hardware_tests/cpu.yaml
+4
@@ -12,6 +12,9 @@ steps:
12 - vllm/_custom_ops.py
13 - tests/kernels/attention/test_cpu_attn.py
14 - tests/v1/attention/test_group_head_counts.py
15 + - tests/kernels/attention/test_amx_mla.py
16 + - vllm/v1/attention/backends/mla/amx_mla.py
17 + - vllm/model_executor/layers/attention/mla_attention.py
18 - tests/kernels/moe/test_cpu_fused_moe.py
19 - tests/kernels/moe/test_cpu_quant_fused_moe.py
20 - tests/kernels/test_onednn.py
@@ -29,6 +32,7 @@ steps:
32 bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
33 pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
34 pytest -x -v -s tests/v1/attention/test_group_head_counts.py
35 + pytest -x -v -s tests/kernels/attention/test_amx_mla.py
36 pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
37 pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
38 pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
cmake/cpu_extension.cmake
+5 -1
@@ -519,7 +519,11 @@ if (ENABLE_X86_ISA)
519 "csrc/cpu/sgl-kernels/moe.cpp"
520 "csrc/cpu/sgl-kernels/moe_int8.cpp"
521 "csrc/cpu/sgl-kernels/moe_int4.cpp"
522 - "csrc/cpu/sgl-kernels/moe_fp8.cpp")
522 + "csrc/cpu/sgl-kernels/moe_fp8.cpp"
523 + "csrc/cpu/sgl-kernels/bmm.cpp"
524 + "csrc/cpu/sgl-kernels/decode.cpp"
525 + "csrc/cpu/sgl-kernels/extend.cpp"
526 + "csrc/cpu/sgl-kernels/mla_cache.cpp")
527
528 set(VLLM_EXT_SRC_AVX512
529 "csrc/cpu/sgl-kernels/fla.cpp"
csrc/cpu/sgl-kernels/bmm.cpp new
+208
@@ -0,0 +1,208 @@
1 +// Adapted from
2 +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3 +
4 +// clang-format off
5 +
6 +#include "common.h"
7 +#include "gemm.h"
8 +#include "vec.h"
9 +
10 +namespace {
11 +
12 +template <typename scalar_t, typename packed_t>
13 +void bmm_kernel_impl(
14 + scalar_t* __restrict__ out,
15 + const scalar_t* __restrict__ mat1,
16 + const packed_t* __restrict__ mat2,
17 + int64_t B,
18 + int64_t M,
19 + int64_t N,
20 + int64_t K,
21 + int64_t mat1_strideB,
22 + int64_t mat1_strideM,
23 + int64_t out_strideB,
24 + int64_t out_strideM,
25 + float scale = 0.f) {
26 + constexpr int64_t BLOCK_M = block_size_m();
27 + constexpr int64_t BLOCK_N = block_size_n();
28 + const int64_t MB = div_up(M, BLOCK_M);
29 + const int64_t NB = div_up(N, BLOCK_N);
30 +
31 + // mat2 contiguous in [B, N, K]
32 + int64_t mat2_strideB = N * K;
33 + int64_t mat2_strideN = K;
34 +
35 + const bool use_brgemm = can_use_brgemm<scalar_t>(M);
36 +
37 + // parallel on [B, MB, NB]
38 + at::parallel_for(0, B * MB * NB, 0, [&](int64_t begin, int64_t end) {
39 + int64_t bs{0}, mb{0}, nb{0};
40 + data_index_init(begin, bs, B, mb, MB, nb, NB);
41 +
42 + // for brgemm, use float32 for accumulate
43 + alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
44 +
45 + for (int i = begin; i < end; ++i) {
46 + UNUSED(i);
47 + int mb_start = mb * BLOCK_M;
48 + int mb_size = std::min(M - mb_start, BLOCK_M);
49 + int nb_start = nb * BLOCK_N;
50 + int nb_size = std::min(N - nb_start, BLOCK_N);
51 +
52 + tinygemm_kernel<scalar_t>(
53 + /* A */ mat1 + bs * mat1_strideB + mb_start * mat1_strideM,
54 + /* B */ mat2 + bs * mat2_strideB + nb_start * mat2_strideN /* nb * BLOCK_N * K */,
55 + /* C */ out + bs * out_strideB + mb_start * out_strideM + nb_start,
56 + /* Ctmp*/ Ctmp,
57 + /* M */ mb_size,
58 + /* N */ nb_size,
59 + /* K */ K,
60 + /* lda */ mat1_strideM,
61 + /* ldb */ nb_size,
62 + /* ldc */ out_strideM,
63 + /* brg */ use_brgemm);
64 +
65 + // move to the next index
66 + data_index_step(bs, B, mb, MB, nb, NB);
67 + }
68 +
69 + if (use_brgemm) {
70 + at::native::cpublas::brgemm_release();
71 + }
72 + });
73 +}
74 +
75 +template <>
76 +void bmm_kernel_impl(
77 + at::BFloat16* __restrict__ out,
78 + const at::BFloat16* __restrict__ mat1,
79 + const at::Float8_e4m3fn* __restrict__ mat2,
80 + int64_t B,
81 + int64_t M,
82 + int64_t N,
83 + int64_t K,
84 + int64_t mat1_strideB,
85 + int64_t mat1_strideM,
86 + int64_t out_strideB,
87 + int64_t out_strideM,
88 + float scale) {
89 + constexpr int64_t BLOCK_M = block_size_m();
90 + constexpr int64_t BLOCK_N = block_size_n();
91 + const int64_t MB = div_up(M, BLOCK_M);
92 + const int64_t NB = div_up(N, BLOCK_N);
93 +
94 + // mat2 contiguous in [B, N, K]
95 + int64_t mat2_strideB = N * K;
96 + int64_t mat2_strideN = K;
97 +
98 + const bool use_brgemm = can_use_brgemm<at::BFloat16>(M);
99 +
100 + // parallel on [B, MB, NB]
101 + parallel_2d(B * MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
102 + // for brgemm, use float32 for accumulate
103 + alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
104 + // for brgemm when mat2 is float8_e4m3
105 + alignas(64) at::BFloat16 Btmp[BLOCK_N * BLOCK_K];
106 +
107 + loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
108 + int64_t bs = mb / MB;
109 + int64_t mb_start = (mb % MB) * BLOCK_M;
110 + int64_t mb_size = std::min(M - mb_start, BLOCK_M);
111 + int64_t nb_start = nb * BLOCK_N;
112 + int64_t nb_size = std::min(N - nb_start, BLOCK_N);
113 +
114 + tinygemm_kernel(
115 + /* A */ mat1 + bs * mat1_strideB + mb_start * mat1_strideM,
116 + /* B */ mat2 + bs * mat2_strideB + nb_start * mat2_strideN /* nb * BLOCK_N * K */,
117 + /* C */ out + bs * out_strideB + mb_start * out_strideM + nb_start,
118 + /* Btmp*/ Btmp,
119 + /* Ctmp*/ Ctmp,
120 + /*scale*/ scale,
121 + /* M */ mb_size,
122 + /* N */ nb_size,
123 + /* K */ K,
124 + /* lda */ mat1_strideM,
125 + /* ldb */ nb_size,
126 + /* ldc */ out_strideM,
127 + /* brg */ use_brgemm);
128 + });
129 +
130 + if (use_brgemm) {
131 + at::native::cpublas::brgemm_release();
132 + }
133 + });
134 +}
135 +
136 +} // anonymous namespace
137 +
138 +// mat1 : [B, M, K]
139 +// mat2 : [B, N, K] or [B, OC, IC]
140 +// out : [B, M, N]
141 +// scale: [] 0-dim tensor for per tensor quant
142 +//
143 +void bmm_cpu(
144 + at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni, const std::optional<at::Tensor>& scale) {
145 + auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
146 +
147 + // input and out could be non-contiguous
148 + // weight needs to be contiguous in [OC, IC] order
149 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
150 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(out);
151 + CHECK_INPUT(mat2);
152 + CHECK_DIM(3, out);
153 + CHECK_DIM(3, mat1);
154 + CHECK_DIM(3, mat2);
155 +
156 + int64_t B = mat1.size(0);
157 + int64_t M = mat1.size(1);
158 + int64_t N = mat2.size(1);
159 + int64_t K = mat1.size(2);
160 +
161 + const bool use_fp8_w8a16 = scale.has_value();
162 + TORCH_CHECK(N % 32 == 0, "tinygemm requires N to be 32x.");
163 +
164 + int64_t mat1_strideB = mat1.stride(0);
165 + int64_t mat1_strideM = mat1.stride(1);
166 + int64_t out_strideB = out.stride(0);
167 + int64_t out_strideM = out.stride(1);
168 +
169 + // check shapes
170 + TORCH_CHECK(mat2.size(0) == B && mat2.size(2) == K, "bmm: mat2 shape mismatch!");
171 + TORCH_CHECK(out.size(0) == B && out.size(1) == M, "bmm: out shape mismatch!");
172 + if (!use_fp8_w8a16) {
173 + AT_DISPATCH_REDUCED_FLOATING_TYPES(mat1.scalar_type(), "bmm_kernel_impl", [&] {
174 + bmm_kernel_impl<scalar_t, scalar_t>(
175 + out.data_ptr<scalar_t>(),
176 + mat1.data_ptr<scalar_t>(),
177 + packed_w.data_ptr<scalar_t>(),
178 + B,
179 + M,
180 + N,
181 + K,
182 + mat1_strideB,
183 + mat1_strideM,
184 + out_strideB,
185 + out_strideM);
186 + });
187 + } else { // fp8 bmm
188 + float scale_val = 0.f;
189 +
190 + auto scale_tensor = scale.value();
191 + TORCH_CHECK(scale_tensor.ndimension() == 0, "bmm: expect scale to be 0-dim tensor.");
192 + scale_val = scale_tensor.item<float>();
193 +
194 + bmm_kernel_impl<at::BFloat16, at::Float8_e4m3fn>(
195 + out.data_ptr<at::BFloat16>(),
196 + mat1.data_ptr<at::BFloat16>(),
197 + packed_w.data_ptr<at::Float8_e4m3fn>(),
198 + B,
199 + M,
200 + N,
201 + K,
202 + mat1_strideB,
203 + mat1_strideM,
204 + out_strideB,
205 + out_strideM,
206 + scale_val);
207 + }
208 +}
csrc/cpu/sgl-kernels/decode.cpp new
+1821
@@ -0,0 +1,1821 @@
1 +// Adapted from
2 +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3 +
4 +// clang-format off
5 +
6 +#include "common.h"
7 +#include "gemm.h"
8 +#include "vec.h"
9 +
10 +namespace {
11 +
12 +// [NOTE] TODO list for this kernel:
13 +// 1. tune the value for BLOCK_N
14 +// 2. planning for {batches, num_heads, num_kv_splits}
15 +// and use actual num_kv_splits for small seq length
16 +// 3. try fast impl of `.tanh()`
17 +// 4. provide amx kernel for index_gemm_kernel_nn when M = 16
18 +//
19 +
20 +#if defined(CPU_CAPABILITY_AVX512)
21 +// key: from [N, 32] to [32/2, N, 2]
22 +// val: from [N, 32] to [N/2, 32, 2]
23 +template <typename scalar_t, typename index_t>
24 +inline void pack_vnni_Nx32(
25 + scalar_t* __restrict__ dst0,
26 + scalar_t* __restrict__ dst1,
27 + const scalar_t* __restrict__ src,
28 + const index_t* __restrict__ ind,
29 + int N,
30 + int ld_src,
31 + int ld_dst0,
32 + int ld_dst1,
33 + bool convert_v) {
34 + __m512i vinputs[16];
35 + int n = 0;
36 + for (; n < N; ++n) {
37 + vinputs[n] = _mm512_loadu_si512(src + ind[n] * ld_src);
38 + }
39 + // padding with zero to avoid uninitialized vectors
40 + for (; n < 16; ++n) {
41 + vinputs[n] = _mm512_set1_epi32(0);
42 + }
43 +
44 + // pack value, skip 64 elems for deepseek
45 + // handle 2 vectors at a time from [2, 32] to [32, 2]
46 + if (convert_v) {
47 + for (int n = 0; n < 16; n += 2) {
48 + __m512i d0, d1;
49 + std::tie(d0, d1) = transpose_2x32_16bit(vinputs[n], vinputs[n + 1]);
50 + _mm512_storeu_si512(dst1 + (n >> 1) * ld_dst1 * 2, d0);
51 + _mm512_storeu_si512(dst1 + (n >> 1) * ld_dst1 * 2 + 32, d1);
52 + }
53 + }
54 +
55 + // pack key
56 + transpose_16x16_32bit(vinputs);
57 +
58 + const __mmask16 vmask = (1 << N) - 1;
59 + for (int k = 0; k < 16; ++k) {
60 + _mm512_mask_storeu_epi32(dst0 + k * ld_dst0 * 2, vmask, vinputs[k]);
61 + }
62 +}
63 +#endif
64 +
65 +// [NOTE]: MLA vnni format conversion
66 +//
67 +// here we apply same strategy as `FlashMLA`:
68 +// each kv_cache is loaded once and packed twice (L2 cache hit)
69 +//
70 +// * for key: from [N, K/2, 2] to [K/2, N, 2]
71 +// * for value: from [N/2, 2, Kv] to [N/2, Kv, 2]
72 +//
73 +template <typename scalar_t, typename index_t>
74 +void pack_vnni(
75 + scalar_t* __restrict__ dst0,
76 + scalar_t* __restrict__ dst1,
77 + const scalar_t* __restrict__ src,
78 + const index_t* __restrict__ ind,
79 + int N,
80 + int K,
81 + int Kv,
82 + int ld_src,
83 + int ld_dst0,
84 + int ld_dst1) {
85 +#if defined(CPU_CAPABILITY_AVX512)
86 + const int NB = div_up(N, 16);
87 + const int KB = K / 32; // no remainder
88 + const int KBv = Kv / 32; // no remainder
89 +
90 + for (int nb = 0; nb < NB; ++nb) {
91 + for (int kb = 0; kb < KB; ++kb) {
92 + // handle 16x512bits each block
93 + int nb_size = std::min(N - nb * 16, 16);
94 + pack_vnni_Nx32<scalar_t, index_t>(
95 + /* dst0 */ dst0 + ((kb * 32) >> 1) * ld_dst0 * 2 + nb * 16 * 2,
96 + /* dst1 */ dst1 + ((nb * 16) >> 1) * ld_dst1 * 2 + kb * 32 * 2,
97 + /* src */ src + kb * 32,
98 + /* ind */ ind + nb * 16,
99 + /* N */ nb_size,
100 + /* ld_src */ ld_src,
101 + /* ld_dst0 */ ld_dst0,
102 + /* ld_dst1 */ ld_dst1,
103 + /* cvt_v */ kb < KBv);
104 + }
105 + }
106 +#else
107 + for (int n = 0; n < N; ++n) {
108 + index_t index = ind[n];
109 + for (int k = 0; k < K / 2; ++k) {
110 + for (int d = 0; d < 2; ++d) {
111 + dst0[k * ld_dst0 * 2 + n * 2 + d] = src[index * ld_src + k * 2 + d];
112 + }
113 + }
114 + }
115 + // from [N/2, 2, K] to [N/2, K, 2]
116 + for (int n = 0; n < (N >> 1) * 2; n += 2) {
117 + index_t index0 = ind[n + 0];
118 + index_t index1 = ind[n + 1];
119 + for (int k = 0; k < Kv; ++k) {
120 + dst1[(n >> 1) * ld_dst1 * 2 + k * 2 + 0] = src[index0 * ld_src + k];
121 + dst1[(n >> 1) * ld_dst1 * 2 + k * 2 + 1] = src[index1 * ld_src + k];
122 + }
123 + }
124 + if (N % 2 != 0) {
125 + index_t index = ind[N - 1];
126 + for (int k = 0; k < Kv; ++k) {
127 + dst1[(N >> 1) * ld_dst1 * 2 + k * 2 + 0] = src[index * ld_src + k];
128 + dst1[(N >> 1) * ld_dst1 * 2 + k * 2 + 1] = 0;
129 + }
130 + }
131 +#endif
132 +}
133 +
134 +template <typename scalar_t>
135 +inline void fill_stub(scalar_t* __restrict__ out, float val, int64_t size) {
136 + using Vec = at::vec::Vectorized<scalar_t>;
137 + constexpr int kVecSize = Vec::size();
138 + const Vec data_vec = Vec(static_cast<scalar_t>(val));
139 + int64_t d = 0;
140 +#pragma GCC unroll 4
141 + for (; d <= size - kVecSize; d += kVecSize) {
142 + data_vec.store(out + d);
143 + }
144 + if (size - d > 0) {
145 + data_vec.store(out + d, size - d);
146 + }
147 +}
148 +
149 +template <typename scalar_t>
150 +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ acc, float s, int64_t size) {
151 + using bVec = at::vec::Vectorized<scalar_t>;
152 + using fVec = at::vec::Vectorized<float>;
153 + constexpr int kVecSize = bVec::size();
154 + const fVec s_fvec = fVec(s);
155 + int64_t d = 0;
156 +#pragma GCC unroll 4
157 + for (; d <= size - kVecSize; d += kVecSize) {
158 + auto [a_fvec0, a_fvec1] = load_float_vec2(acc + d);
159 + a_fvec0 = a_fvec0 * s_fvec;
160 + a_fvec1 = a_fvec1 * s_fvec;
161 + bVec out_bvec = convert_from_float_ext<scalar_t>(a_fvec0, a_fvec1);
162 + out_bvec.store(out + d);
163 + }
164 + for (; d < size; ++d) {
165 + out[d] = static_cast<scalar_t>(acc[d] * s);
166 + }
167 +}
168 +
169 +template <typename scalar_t>
170 +inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ src, int64_t size) {
171 + using bVec = at::vec::Vectorized<scalar_t>;
172 + constexpr int kVecSize = bVec::size();
173 + int64_t d = 0;
174 +#pragma GCC unroll 4
175 + for (; d <= size - kVecSize; d += kVecSize) {
176 + bVec out_bvec = bVec::loadu(src + d);
177 + out_bvec.store(out + d);
178 + }
179 + for (; d < size; ++d) {
180 + out[d] = src[d];
181 + }
182 +}
183 +
184 +template <typename scalar_t, int BLOCK_N>
185 +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input) {
186 + static_assert(BLOCK_N % 32 == 0);
187 + using bVec = at::vec::Vectorized<scalar_t>;
188 + using fVec = at::vec::Vectorized<float>;
189 +
190 + constexpr int COLS = BLOCK_N / 16;
191 + auto store = [&](auto i) {
192 + constexpr int col = i % COLS;
193 + // for COLS = 2, 4 use 512bit store
194 + if constexpr (col % 2 == 0) {
195 + auto [a_fvec0, a_fvec1] = load_float_vec2(input + col * 16);
196 + bVec out_bvec = convert_from_float_ext<scalar_t>(a_fvec0, a_fvec1);
197 + out_bvec.store(out + col * 16);
198 + }
199 + };
200 + Unroll<COLS>{}(store);
201 +}
202 +
203 +// GEMM handles query @ key (indexed) x scale
204 +// A : [M, K]
205 +// B : [N, K] indexed
206 +// C : [M, N]
207 +//
208 +template <typename scalar_t, typename index_t, int BLOCK_M, int BLOCK_N>
209 +struct tinygemm_kernel_nt {
210 + static inline void apply(
211 + const scalar_t* __restrict__ A,
212 + const scalar_t* __restrict__ B,
213 + float* __restrict__ C,
214 + const index_t* __restrict__ indices,
215 + float scale,
216 + int64_t lda,
217 + int64_t ldb,
218 + int64_t ldc,
219 + int64_t K,
220 + int64_t max_tokens) {
221 + for (int64_t m = 0; m < BLOCK_M; ++m) {
222 + for (int64_t n = 0; n < BLOCK_N; ++n) {
223 + float sum = 0.f;
224 + int64_t b_idx = indices[n];
225 + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!");
226 + for (int64_t k = 0; k < K; ++k) {
227 + sum += scale * static_cast<float>(A[m * lda + k]) * static_cast<float>(B[b_idx * ldb + k]);
228 + }
229 + C[m * ldc + n] = sum;
230 + }
231 + }
232 + }
233 +};
234 +
235 +#if defined(CPU_CAPABILITY_AVX512)
236 +template <typename index_t, int BLOCK_M, int BLOCK_N>
237 +struct tinygemm_kernel_nt<at::BFloat16, index_t, BLOCK_M, BLOCK_N> {
238 + static inline void apply(
239 + const at::BFloat16* __restrict__ A,
240 + const at::BFloat16* __restrict__ B,
241 + float* __restrict__ C,
242 + const index_t* __restrict__ indices,
243 + float scale,
244 + int64_t lda,
245 + int64_t ldb,
246 + int64_t ldc,
247 + int64_t K,
248 + int64_t max_tokens) {
249 + constexpr int ROWS = BLOCK_M;
250 + constexpr int COLS = BLOCK_N;
251 +
252 + __m512bh va;
253 + __m512bh vb[COLS];
254 + __m512 vc[ROWS * COLS];
255 + __m512 vscale = _mm512_set1_ps(scale);
256 +
257 + auto loadc = [&](auto i) { vc[i] = _mm512_setzero_ps(); };
258 + Unroll<ROWS * COLS>{}(loadc);
259 +
260 + // for main loop
261 + auto compute = [&](auto i, int64_t k) {
262 + constexpr int row = i / COLS;
263 + constexpr int col = i % COLS;
264 +
265 + if constexpr (col == 0) {
266 + va = (__m512bh)(_mm512_loadu_si512(A + row * lda + k));
267 + }
268 + if constexpr (row == 0) {
269 + if constexpr (col + 1 < COLS) {
270 + int64_t b_idx_prefetch = indices[col + 1];
271 + _mm_prefetch(B + b_idx_prefetch * ldb + k, _MM_HINT_T0);
272 + }
273 + int64_t b_idx = indices[col];
274 + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!");
275 + vb[col] = (__m512bh)(_mm512_loadu_si512(B + b_idx * ldb + k));
276 + }
277 + vc[i] = _mm512_dpbf16_ps(vc[i], va, vb[col]);
278 + };
279 +
280 + // for remainder
281 + auto compute2 = [&](auto i, int64_t k, __mmask32 mask) {
282 + constexpr int row = i / COLS;
283 + constexpr int col = i % COLS;
284 +
285 + if constexpr (col == 0) {
286 + va = (__m512bh)(_mm512_maskz_loadu_epi16(mask, A + row * lda + k));
287 + }
288 + if constexpr (row == 0) {
289 + int64_t b_idx = indices[col];
290 + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!");
291 + vb[col] = (__m512bh)(_mm512_maskz_loadu_epi16(mask, B + b_idx * ldb + k));
292 + }
293 + vc[i] = _mm512_dpbf16_ps(vc[i], va, vb[col]);
294 + };
295 +
296 + int64_t k = 0;
297 + for (; k <= K - 32; k += 32) {
298 + Unroll<ROWS * COLS>{}(compute, k);
299 + }
300 + int64_t count = K - k;
301 + if (count > 0) {
302 + __mmask32 mask = (1ULL << count) - 1;
303 + Unroll<ROWS * COLS>{}(compute2, k, mask);
304 + }
305 +
306 + auto storec = [&](auto i) {
307 + constexpr int row = i / COLS;
308 + constexpr int col = i % COLS;
309 + C[row * ldc + col] = _mm512_reduce_add_ps(_mm512_mul_ps(vc[i], vscale));
310 + };
311 + Unroll<ROWS * COLS>{}(storec);
312 + }
313 +};
314 +#endif
315 +
316 +#if defined(CPU_CAPABILITY_AVX512)
317 +template <typename index_t, int BLOCK_M, int BLOCK_N>
318 +struct tinygemm_kernel_nt<at::Half, index_t, BLOCK_M, BLOCK_N> {
319 + static inline void apply(
320 + const at::Half* __restrict__ A,
321 + const at::Half* __restrict__ B,
322 + float* __restrict__ C,
323 + const index_t* __restrict__ indices,
324 + float scale,
325 + int64_t lda,
326 + int64_t ldb,
327 + int64_t ldc,
328 + int64_t K,
329 + int64_t max_tokens) {
330 + constexpr int ROWS = BLOCK_M;
331 + constexpr int COLS = BLOCK_N;
332 +
333 + __m512 va0, va1;
334 + __m512 vb0[COLS], vb1[COLS];
335 + __m512 vc[ROWS * COLS];
336 + __m512 vscale = _mm512_set1_ps(scale);
337 +
338 + auto loadc = [&](auto i) { vc[i] = _mm512_setzero_ps(); };
339 + Unroll<ROWS * COLS>{}(loadc);
340 +
341 + auto compute = [&](auto i, int64_t k) {
342 + constexpr int row = i / COLS;
343 + constexpr int col = i % COLS;
344 +
345 + if constexpr (col == 0) {
346 + __m512i a16 = _mm512_loadu_si512((__m512i const*)(A + row * lda + k));
347 + va0 = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(a16, 0));
348 + va1 = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(a16, 1));
349 + }
350 +
351 + if constexpr (row == 0) {
352 + int64_t b_idx = indices[col];
353 + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!");
354 + __m512i b16 = _mm512_loadu_si512((__m512i const*)(B + b_idx * ldb + k));
355 + vb0[col] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0));
356 + vb1[col] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1));
357 + }
358 +
359 + vc[i] = _mm512_fmadd_ps(va0, vb0[col], _mm512_fmadd_ps(va1, vb1[col], vc[i]));
360 + };
361 +
362 + auto compute2 = [&](auto i, int64_t k, __mmask32 mask) {
363 + constexpr int row = i / COLS;
364 + constexpr int col = i % COLS;
365 +
366 + if constexpr (col == 0) {
367 + __m512i a16 = _mm512_maskz_loadu_epi16(mask, (const void*)(A + row * lda + k));
368 + va0 = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(a16, 0));
369 + va1 = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(a16, 1));
370 + }
371 +
372 + if constexpr (row == 0) {
373 + int64_t b_idx = indices[col];
374 + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!");
375 + __m512i b16 = _mm512_maskz_loadu_epi16(mask, (const void*)(B + b_idx * ldb + k));
376 + vb0[col] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0));
377 + vb1[col] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1));
378 + }
379 +
380 + vc[i] = _mm512_fmadd_ps(va0, vb0[col], _mm512_fmadd_ps(va1, vb1[col], vc[i]));
381 + };
382 +
383 + int64_t k = 0;
384 + for (; k <= K - 32; k += 32) {
385 + Unroll<ROWS * COLS>{}(compute, k);
386 + }
387 + int64_t count = K - k;
388 + if (count > 0) {
389 + __mmask32 mask = (1ULL << count) - 1;
390 + Unroll<ROWS * COLS>{}(compute2, k, mask);
391 + }
392 +
393 + auto storec = [&](auto i) {
394 + constexpr int row = i / COLS;
395 + constexpr int col = i % COLS;
396 + C[row * ldc + col] = _mm512_reduce_add_ps(_mm512_mul_ps(vc[i], vscale));
397 + };
398 + Unroll<ROWS * COLS>{}(storec);
399 + }
400 +};
401 +#endif
402 +
403 +#define LAUNCH_TINYGEMM_KERNEL_NT(MB_SIZE, NB_SIZE) \
404 + tinygemm_kernel_nt<scalar_t, index_t, MB_SIZE, NB_SIZE>::apply( \
405 + A + mb_start * lda, B, C + mb_start * ldc + nb_start, indices + nb_start, scale, lda, ldb, ldc, K, max_tokens);
406 +
407 +// this is used when N isn't multiple of 16,
408 +// N corresponds to `head_size_v` which should be 16x
409 +template <typename scalar_t, typename index_t>
410 +inline void tinygemm_kernel_nn_scalar(
411 + const float* __restrict__ A,
412 + const scalar_t* __restrict__ B,
413 + float* __restrict__ C,
414 + const index_t* __restrict__ indices,
415 + const float* __restrict__ scale,
416 + int64_t M,
417 + int64_t N,
418 + int64_t K,
419 + int64_t lda,
420 + int64_t ldb,
421 + int64_t ldc,
422 + int64_t max_tokens) {
423 + for (int64_t m = 0; m < M; ++m) {
424 + for (int64_t n = 0; n < N; ++n) {
425 + C[m * ldc + n] *= scale[m];
426 + for (int64_t k = 0; k < K; ++k) {
427 + int64_t b_idx = indices[k];
428 + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!");
429 + C[m * ldc + n] += A[m * lda + k] * static_cast<float>(B[b_idx * ldb + n]);
430 + }
431 + }
432 + }
433 +}
434 +
435 +// GEMM handles v' * scale + attn @ value (indexed)
436 +// A : [M, K]
437 +// B : [K, N] indexed
438 +// C ：[M, N]
439 +//
440 +template <typename scalar_t, typename index_t, int BLOCK_M, int BLOCK_N>
441 +struct tinygemm_kernel_nn {
442 + static inline void apply(
443 + const float* __restrict__ A,
444 + const scalar_t* __restrict__ B,
445 + float* __restrict__ C,
446 + const index_t* __restrict__ indices,
447 + const float* __restrict__ scale,
448 + int64_t lda,
449 + int64_t ldb,
450 + int64_t ldc,
451 + int64_t K,
452 + int64_t max_tokens) {
453 + tinygemm_kernel_nn_scalar(A, B, C, indices, scale, BLOCK_M, BLOCK_N, K, lda, ldb, ldc, max_tokens);
454 + }
455 +};
456 +
457 +#if defined(CPU_CAPABILITY_AVX512)
458 +template <typename index_t, int BLOCK_M, int BLOCK_N>
459 +struct tinygemm_kernel_nn<at::BFloat16, index_t, BLOCK_M, BLOCK_N> {
460 + static inline void apply(
461 + const float* __restrict__ A,
462 + const at::BFloat16* __restrict__ B,
463 + float* __restrict__ C,
464 + const index_t* __restrict__ indices,
465 + const float* __restrict__ scale,
466 + int64_t lda,
467 + int64_t ldb,
468 + int64_t ldc,
469 + int64_t K,
470 + int64_t max_tokens) {
471 + constexpr int ROWS = BLOCK_M;
472 + constexpr int COLS = BLOCK_N / 16;
473 +
474 + __m512 va;
475 + __m512 vb[COLS];
476 + __m512 vc[ROWS * COLS];
477 + __m512 vscale;
478 +
479 + auto loadc = [&](auto i) {
480 + constexpr int row = i / COLS;
481 + constexpr int col = i % COLS;
482 +#pragma GCC diagnostic push
483 +#pragma GCC diagnostic ignored "-Warray-bounds"
484 + if constexpr (col == 0) {
485 + vscale = _mm512_set1_ps(scale[row]);
486 + }
487 +#pragma GCC diagnostic pop
488 + vc[i] = _mm512_loadu_ps(C + row * ldc + col * 16);
489 + vc[i] = _mm512_mul_ps(vc[i], vscale);
490 + };
491 + Unroll<ROWS * COLS>{}(loadc);
492 +
493 + auto compute = [&](auto i, int64_t k) {
494 + constexpr int row = i / COLS;
495 + constexpr int col = i % COLS;
496 +
497 + if constexpr (col == 0) {
498 + va = _mm512_set1_ps(A[row * lda + k]);
499 + }
500 + if constexpr (row == 0) {
501 + if (k + 1 < K) {
502 + int64_t b_idx_prefetch = indices[k + 1];
503 + _mm_prefetch(B + b_idx_prefetch * ldb + col * 16, _MM_HINT_T0);
504 + }
505 + int64_t b_idx = indices[k];
506 + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!");
507 +
508 + // for COLS = 2, 4, 6, 8 use 512 bit load
509 + // for COLS = 1, 3, 5, 7 use 256 bit load
510 + if constexpr (COLS % 2 == 0) {
511 + if constexpr (col % 2 == 0) {
512 + __m512i b16 = _mm512_loadu_si512(reinterpret_cast<const __m512i*>(B + b_idx * ldb + col * 16));
513 + vb[col + 0] = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0));
514 + vb[col + 1] = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1));
515 + }
516 + } else {
517 + __m256i b16 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(B + b_idx * ldb + col * 16));
518 + vb[col] = CVT_BF16_TO_FP32(b16);
519 + }
520 + }
521 + vc[i] = _mm512_fmadd_ps(va, vb[col], vc[i]);
522 + };
523 +
524 + for (int64_t k = 0; k < K; ++k) {
525 + Unroll<ROWS * COLS>{}(compute, k);
526 + }
527 +
528 + auto storec = [&](auto i) {
529 + constexpr int row = i / COLS;
530 + constexpr int col = i % COLS;
531 + _mm512_storeu_ps(C + row * ldc + col * 16, vc[i]);
532 + };
533 + Unroll<ROWS * COLS>{}(storec);
534 + }
535 +};
536 +#endif
537 +
538 +#if defined(CPU_CAPABILITY_AVX512)
539 +template <typename index_t, int BLOCK_M, int BLOCK_N>
540 +struct tinygemm_kernel_nn<at::Half, index_t, BLOCK_M, BLOCK_N> {
541 + static inline void apply(
542 + const float* __restrict__ A,
543 + const at::Half* __restrict__ B,
544 + float* __restrict__ C,
545 + const index_t* __restrict__ indices,
546 + const float* __restrict__ scale,
547 + int64_t lda,
548 + int64_t ldb,
549 + int64_t ldc,
550 + int64_t K,
551 + int64_t max_tokens) {
552 + constexpr int ROWS = BLOCK_M;
553 + constexpr int COLS = BLOCK_N / 16;
554 +
555 + __m512 va;
556 + __m512 vb[COLS];
557 + __m512 vc[ROWS * COLS];
558 + __m512 vscale;
559 +
560 + auto loadc = [&](auto i) {
561 + constexpr int row = i / COLS;
562 + constexpr int col = i % COLS;
563 +#pragma GCC diagnostic push
564 +#pragma GCC diagnostic ignored "-Warray-bounds"
565 + if constexpr (col == 0) {
566 + vscale = _mm512_set1_ps(scale[row]);
567 + }
568 +#pragma GCC diagnostic pop
569 + vc[i] = _mm512_loadu_ps(C + row * ldc + col * 16);
570 + vc[i] = _mm512_mul_ps(vc[i], vscale);
571 + };
572 + Unroll<ROWS * COLS>{}(loadc);
573 +
574 + auto compute = [&](auto i, int64_t k) {
575 + constexpr int row = i / COLS;
576 + constexpr int col = i % COLS;
577 +
578 + if constexpr (col == 0) {
579 + va = _mm512_set1_ps(A[row * lda + k]);
580 + }
581 + if constexpr (row == 0) {
582 + if (k + 1 < K) {
583 + int64_t b_idx_prefetch = indices[k + 1];
584 + _mm_prefetch(B + b_idx_prefetch * ldb + col * 16, _MM_HINT_T0);
585 + }
586 + int64_t b_idx = indices[k];
587 + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!");
588 +
589 + // for COLS = 2, 4, 6, 8 use 512 bit load
590 + // for COLS = 1, 3, 5, 7 use 256 bit load
591 + if constexpr (COLS % 2 == 0) {
592 + if constexpr (col % 2 == 0) {
593 + __m512i b16 = _mm512_loadu_si512(reinterpret_cast<const __m512i*>(B + b_idx * ldb + col * 16));
594 + vb[col + 0] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0));
595 + vb[col + 1] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1));
596 + }
597 + } else {
598 + __m256i b16 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(B + b_idx * ldb + col * 16));
599 + vb[col] = CVT_FP16_TO_FP32(b16);
600 + }
601 + }
602 + vc[i] = _mm512_fmadd_ps(va, vb[col], vc[i]);
603 + };
604 +
605 + for (int64_t k = 0; k < K; ++k) {
606 + Unroll<ROWS * COLS>{}(compute, k);
607 + }
608 +
609 + auto storec = [&](auto i) {
610 + constexpr int row = i / COLS;
611 + constexpr int col = i % COLS;
612 + _mm512_storeu_ps(C + row * ldc + col * 16, vc[i]);
613 + };
614 + Unroll<ROWS * COLS>{}(storec);
615 + }
616 +};
617 +#endif
618 +
619 +#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE) \
620 + tinygemm_kernel_nn<scalar_t, index_t, MB_SIZE, NB_SIZE>::apply( \
621 + A + mb_start * lda, \
622 + B + nb_start, \
623 + C + mb_start * ldc + nb_start, \
624 + indices, \
625 + scale + mb_start, \
626 + lda, \
627 + ldb, \
628 + ldc, \
629 + K, \
630 + max_tokens);
631 +
632 +template <typename scalar_t, typename index_t>
633 +void index_gemm_kernel_nt(
634 + const scalar_t* __restrict__ A,
635 + const scalar_t* __restrict__ B,
636 + float* __restrict__ C,
637 + const index_t* __restrict__ indices,
638 + float scale,
639 + int64_t M,
640 + int64_t N,
641 + int64_t K,
642 + int64_t lda,
643 + int64_t ldb,
644 + int64_t ldc,
645 + int64_t max_tokens) {
646 + // pattern: 1-8-8
647 + if (M == 1) {
648 + constexpr int64_t BLOCK_N = 8;
649 + const int64_t NB = div_up(N, BLOCK_N);
650 + int64_t mb_start = 0, lda = 1, ldc = 1;
651 +
652 + for (int64_t nb = 0; nb < NB; ++nb) {
653 + int64_t nb_start = nb * BLOCK_N;
654 + int64_t nb_size = std::min(BLOCK_N, N - nb_start);
655 +
656 + switch (nb_size) {
657 + case 1:
658 + LAUNCH_TINYGEMM_KERNEL_NT(1, 1);
659 + break;
660 + case 2:
661 + LAUNCH_TINYGEMM_KERNEL_NT(1, 2);
662 + break;
663 + case 3:
664 + LAUNCH_TINYGEMM_KERNEL_NT(1, 3);
665 + break;
666 + case 4:
667 + LAUNCH_TINYGEMM_KERNEL_NT(1, 4);
668 + break;
669 + case 5:
670 + LAUNCH_TINYGEMM_KERNEL_NT(1, 5);
671 + break;
672 + case 6:
673 + LAUNCH_TINYGEMM_KERNEL_NT(1, 6);
674 + break;
675 + case 7:
676 + LAUNCH_TINYGEMM_KERNEL_NT(1, 7);
677 + break;
678 + case 8:
679 + LAUNCH_TINYGEMM_KERNEL_NT(1, 8);
680 + break;
681 + default:
682 + TORCH_CHECK(false, "Unexpected block size, 1x", "nb_size");
683 + }
684 + }
685 + return;
686 + }
687 +
688 + // default pattern: 1-6-24
689 + // FP16 pattern: 2-8-16
690 + constexpr int64_t BLOCK_M = 4;
691 + constexpr int64_t BLOCK_N = std::is_same_v<scalar_t, at::Half> ? 4 : 6;
692 + const int64_t MB = div_up(M, BLOCK_M);
693 + const int64_t NB = div_up(N, BLOCK_N);
694 +
695 + for (int64_t mb = 0; mb < MB; ++mb) {
696 + int64_t mb_start = mb * BLOCK_M;
697 + int64_t mb_size = std::min(BLOCK_M, M - mb_start);
698 + for (int64_t nb = 0; nb < NB; ++nb) {
699 + int64_t nb_start = nb * BLOCK_N;
700 + int64_t nb_size = std::min(BLOCK_N, N - nb_start);
701 +
702 + switch (mb_size << 4 | nb_size) {
703 + // mb_size = 1
704 + case 0x11:
705 + LAUNCH_TINYGEMM_KERNEL_NT(1, 1);
706 + break;
707 + case 0x12:
708 + LAUNCH_TINYGEMM_KERNEL_NT(1, 2);
709 + break;
710 + case 0x13:
711 + LAUNCH_TINYGEMM_KERNEL_NT(1, 3);
712 + break;
713 + case 0x14:
714 + LAUNCH_TINYGEMM_KERNEL_NT(1, 4);
715 + break;
716 + case 0x15:
717 + LAUNCH_TINYGEMM_KERNEL_NT(1, 5);
718 + break;
719 + case 0x16:
720 + LAUNCH_TINYGEMM_KERNEL_NT(1, 6);
721 + break;
722 + // mb_size = 2
723 + case 0x21:
724 + LAUNCH_TINYGEMM_KERNEL_NT(2, 1);
725 + break;
726 + case 0x22:
727 + LAUNCH_TINYGEMM_KERNEL_NT(2, 2);
728 + break;
729 + case 0x23:
730 + LAUNCH_TINYGEMM_KERNEL_NT(2, 3);
731 + break;
732 + case 0x24:
733 + LAUNCH_TINYGEMM_KERNEL_NT(2, 4);
734 + break;
735 + case 0x25:
736 + LAUNCH_TINYGEMM_KERNEL_NT(2, 5);
737 + break;
738 + case 0x26:
739 + LAUNCH_TINYGEMM_KERNEL_NT(2, 6);
740 + break;
741 + // mb_size = 3
742 + case 0x31:
743 + LAUNCH_TINYGEMM_KERNEL_NT(3, 1);
744 + break;
745 + case 0x32:
746 + LAUNCH_TINYGEMM_KERNEL_NT(3, 2);
747 + break;
748 + case 0x33:
749 + LAUNCH_TINYGEMM_KERNEL_NT(3, 3);
750 + break;
751 + case 0x34:
752 + LAUNCH_TINYGEMM_KERNEL_NT(3, 4);
753 + break;
754 + case 0x35:
755 + LAUNCH_TINYGEMM_KERNEL_NT(3, 5);
756 + break;
757 + case 0x36:
758 + LAUNCH_TINYGEMM_KERNEL_NT(3, 6);
759 + break;
760 + // mb_size = 4
761 + case 0x41:
762 + LAUNCH_TINYGEMM_KERNEL_NT(4, 1);
763 + break;
764 + case 0x42:
765 + LAUNCH_TINYGEMM_KERNEL_NT(4, 2);
766 + break;
767 + case 0x43:
768 + LAUNCH_TINYGEMM_KERNEL_NT(4, 3);
769 + break;
770 + case 0x44:
771 + LAUNCH_TINYGEMM_KERNEL_NT(4, 4);
772 + break;
773 + case 0x45:
774 + LAUNCH_TINYGEMM_KERNEL_NT(4, 5);
775 + break;
776 + case 0x46:
777 + LAUNCH_TINYGEMM_KERNEL_NT(4, 6);
778 + break;
779 + default:
780 + TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
781 + }
782 + }
783 + }
784 +}
785 +
786 +template <typename scalar_t, typename index_t>
787 +void index_gemm_kernel_nn(
788 + const float* __restrict__ A,
789 + const scalar_t* __restrict__ B,
790 + float* __restrict__ C,
791 + const index_t* __restrict__ indices,
792 + float* __restrict__ scale,
793 + int64_t M,
794 + int64_t N,
795 + int64_t K,
796 + int64_t lda,
797 + int64_t ldb,
798 + int64_t ldc,
799 + int64_t max_tokens) {
800 + constexpr int kVecSize = 16;
801 + if ((N & (kVecSize - 1)) != 0) {
802 + tinygemm_kernel_nn_scalar(A, B, C, indices, scale, M, N, K, lda, ldb, ldc, max_tokens);
803 + return;
804 + }
805 +
806 + // pattern: 1-8-8
807 + if (M == 1) {
808 + constexpr int64_t BLOCK_N = 8 * kVecSize;
809 + const int64_t NB = div_up(N, BLOCK_N);
810 + int64_t mb_start = 0, lda = 1, ldc = 1;
811 +
812 + for (int64_t nb = 0; nb < NB; ++nb) {
813 + int64_t nb_start = nb * BLOCK_N;
814 + int64_t nb_size = std::min(BLOCK_N, N - nb_start);
815 +
816 + switch (nb_size >> 4) {
817 + case 1:
818 + LAUNCH_TINYGEMM_KERNEL_NN(1, 16);
819 + break;
820 + case 2:
821 + LAUNCH_TINYGEMM_KERNEL_NN(1, 32);
822 + break;
823 + case 3:
824 + LAUNCH_TINYGEMM_KERNEL_NN(1, 48);
825 + break;
826 + case 4:
827 + LAUNCH_TINYGEMM_KERNEL_NN(1, 64);
828 + break;
829 + case 5:
830 + LAUNCH_TINYGEMM_KERNEL_NN(1, 80);
831 + break;
832 + case 6:
833 + LAUNCH_TINYGEMM_KERNEL_NN(1, 96);
834 + break;
835 + case 7:
836 + LAUNCH_TINYGEMM_KERNEL_NN(1, 112);
837 + break;
838 + case 8:
839 + LAUNCH_TINYGEMM_KERNEL_NN(1, 128);
840 + break;
841 + default:
842 + TORCH_CHECK(false, "Unexpected block size, 1x", "nb_size");
843 + }
844 + }
845 + return;
846 + }
847 +
848 + constexpr int64_t BLOCK_M = 4;
849 + constexpr int64_t BLOCK_N = 6 * kVecSize;
850 + const int64_t MB = div_up(M, BLOCK_M);
851 + const int64_t NB = div_up(N, BLOCK_N);
852 +
853 + for (int64_t mb = 0; mb < MB; ++mb) {
854 + int64_t mb_start = mb * BLOCK_M;
855 + int64_t mb_size = std::min(BLOCK_M, M - mb_start);
856 + for (int64_t nb = 0; nb < NB; ++nb) {
857 + int64_t nb_start = nb * BLOCK_N;
858 + int64_t nb_size = std::min(BLOCK_N, N - nb_start);
859 +
860 + switch (mb_size << 4 | nb_size >> 4) {
861 + // mb_size = 1
862 + case 0x11:
863 + LAUNCH_TINYGEMM_KERNEL_NN(1, 16);
864 + break;
865 + case 0x12:
866 + LAUNCH_TINYGEMM_KERNEL_NN(1, 32);
867 + break;
868 + case 0x13:
869 + LAUNCH_TINYGEMM_KERNEL_NN(1, 48);
870 + break;
871 + case 0x14:
872 + LAUNCH_TINYGEMM_KERNEL_NN(1, 64);
873 + break;
874 + case 0x15:
875 + LAUNCH_TINYGEMM_KERNEL_NN(1, 80);
876 + break;
877 + case 0x16:
878 + LAUNCH_TINYGEMM_KERNEL_NN(1, 96);
879 + break;
880 + // mb_size = 2
881 + case 0x21:
882 + LAUNCH_TINYGEMM_KERNEL_NN(2, 16);
883 + break;
884 + case 0x22:
885 + LAUNCH_TINYGEMM_KERNEL_NN(2, 32);
886 + break;
887 + case 0x23:
888 + LAUNCH_TINYGEMM_KERNEL_NN(2, 48);
889 + break;
890 + case 0x24:
891 + LAUNCH_TINYGEMM_KERNEL_NN(2, 64);
892 + break;
893 + case 0x25:
894 + LAUNCH_TINYGEMM_KERNEL_NN(2, 80);
895 + break;
896 + case 0x26:
897 + LAUNCH_TINYGEMM_KERNEL_NN(2, 96);
898 + break;
899 + // mb_size = 3
900 + case 0x31:
901 + LAUNCH_TINYGEMM_KERNEL_NN(3, 16);
902 + break;
903 + case 0x32:
904 + LAUNCH_TINYGEMM_KERNEL_NN(3, 32);
905 + break;
906 + case 0x33:
907 + LAUNCH_TINYGEMM_KERNEL_NN(3, 48);
908 + break;
909 + case 0x34:
910 + LAUNCH_TINYGEMM_KERNEL_NN(3, 64);
911 + break;
912 + case 0x35:
913 + LAUNCH_TINYGEMM_KERNEL_NN(3, 80);
914 + break;
915 + case 0x36:
916 + LAUNCH_TINYGEMM_KERNEL_NN(3, 96);
917 + break;
918 + // mb_size = 4
919 + case 0x41:
920 + LAUNCH_TINYGEMM_KERNEL_NN(4, 16);
921 + break;
922 + case 0x42:
923 + LAUNCH_TINYGEMM_KERNEL_NN(4, 32);
924 + break;
925 + case 0x43:
926 + LAUNCH_TINYGEMM_KERNEL_NN(4, 48);
927 + break;
928 + case 0x44:
929 + LAUNCH_TINYGEMM_KERNEL_NN(4, 64);
930 + break;
931 + case 0x45:
932 + LAUNCH_TINYGEMM_KERNEL_NN(4, 80);
933 + break;
934 + case 0x46:
935 + LAUNCH_TINYGEMM_KERNEL_NN(4, 96);
936 + break;
937 + default:
938 + TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
939 + }
940 + }
941 + }
942 +}
943 +
944 +template <typename scalar_t>
945 +void decode_set_kv_buffer(
946 + scalar_t* __restrict__ k_buffer,
947 + scalar_t* __restrict__ v_buffer,
948 + const scalar_t* __restrict__ key,
949 + const scalar_t* __restrict__ value,
950 + const int64_t* __restrict__ loc,
951 + int64_t batches,
952 + int64_t num_heads_kv,
953 + int64_t head_size,
954 + int64_t head_size_v,
955 + int64_t k_strideN,
956 + int64_t k_strideH,
957 + int64_t v_strideN,
958 + int64_t v_strideH,
959 + int64_t nk_strideN,
960 + int64_t nk_strideH,
961 + int64_t nv_strideN,
962 + int64_t nv_strideH,
963 + bool is_mla) {
964 + at::parallel_for(0, batches * num_heads_kv, 0, [&](int64_t begin, int64_t end) {
965 + int64_t bs{0}, head_kv_id{0};
966 + data_index_init(begin, bs, batches, head_kv_id, num_heads_kv);
967 +
968 + for (int64_t i = begin; i < end; i++) {
969 + int64_t loc_val = loc[bs];
970 + scalar_t* k_buffer_ptr = k_buffer + loc_val * k_strideN + head_kv_id * k_strideH;
971 + const scalar_t* new_key_ptr = key + bs * nk_strideN + head_kv_id * nk_strideH;
972 + copy_stub<scalar_t>(k_buffer_ptr, new_key_ptr, head_size);
973 + if (!is_mla) {
974 + scalar_t* v_buffer_ptr = v_buffer + loc_val * v_strideN + head_kv_id * v_strideH;
975 + const scalar_t* new_value_ptr = value + bs * nv_strideN + head_kv_id * nv_strideH;
976 + copy_stub<scalar_t>(v_buffer_ptr, new_value_ptr, head_size_v);
977 + }
978 +
979 + // move to the next index
980 + data_index_step(bs, batches, head_kv_id, num_heads_kv);
981 + }
982 + });
983 +}
984 +
985 +template <typename scalar_t>
986 +void decode_accumulate_kv_splits(
987 + scalar_t* __restrict__ output,
988 + float* __restrict__ attn_logits,
989 + const scalar_t* __restrict__ sinks_ptr,
990 + int64_t batches,
991 + int64_t num_heads,
992 + int64_t head_size_v,
993 + int64_t num_kv_splits,
994 + int64_t l_stride1,
995 + int64_t l_stride2,
996 + bool has_sink) {
997 + using Vec = at::vec::Vectorized<float>;
998 +
999 + // parallel on [batches, num_heads]
1000 + at::parallel_for(0, batches * num_heads, 0, [&](int64_t begin, int64_t end) {
1001 + int64_t bi{0}, ni{0};
1002 + data_index_init(begin, bi, batches, ni, num_heads);
1003 + // NB: here we use logits[b][h][0] as acc, since
1004 + // for the first kv split (kv_id == 0):
1005 + // m_delta = std::exp(-inf) = 0
1006 + // e_logic = std::exp(0) = 1
1007 + // acc = acc * m_delta + tv * e_logic = tv
1008 + for (int64_t i = begin; i < end; ++i) {
1009 + float* __restrict__ acc = attn_logits + i * l_stride1;
1010 +
1011 + float s_prime = 0.f;
1012 + float m_prime = -std::numeric_limits<scalar_t>::infinity();
1013 +
1014 + // update acc with from each kv_split
1015 + for (int64_t kv_id = 0; kv_id < num_kv_splits; ++kv_id) {
1016 + float* __restrict__ tv = acc + kv_id * l_stride2;
1017 + const float tlogic = (acc + kv_id * l_stride2)[head_size_v];
1018 +
1019 + float m_i = std::max(tlogic, m_prime);
1020 + float m_delta = std::exp(m_prime - m_i);
1021 + float e_logic = std::exp(tlogic - m_i);
1022 + if (kv_id != 0) {
1023 + at::vec::map2<float>(
1024 + [m_delta, e_logic](Vec x, Vec y) { return x * Vec(m_delta) + y * Vec(e_logic); },
1025 + acc,
1026 + acc,
1027 + tv,
1028 + head_size_v);
1029 + }
1030 +
1031 + s_prime = s_prime * m_delta + e_logic;
1032 + m_prime = m_i;
1033 + }
1034 + if (has_sink) {
1035 + s_prime += std::exp(sinks_ptr[ni] - m_prime);
1036 + }
1037 + copy_stub<scalar_t>(output + i * head_size_v, acc, 1 / s_prime, head_size_v);
1038 + // move to the next index
1039 + data_index_step(bi, batches, ni, num_heads);
1040 + }
1041 + });
1042 +}
1043 +
1044 +template <typename scalar_t, typename index_t, int64_t BLOCK_N>
1045 +void decode_attention_kernel_impl(
1046 + scalar_t* __restrict__ output,
1047 + float* __restrict__ attn_logits,
1048 + const scalar_t* __restrict__ query,
1049 + const scalar_t* __restrict__ k_buffer,
1050 + const scalar_t* __restrict__ v_buffer,
1051 + const index_t* __restrict__ req_to_token,
1052 + const int64_t* __restrict__ req_pool_indices,
1053 + const int64_t* __restrict__ seq_lens,
1054 + const int64_t* __restrict__ encoder_lens,
1055 + const scalar_t* __restrict__ sinks,
1056 + int64_t batches,
1057 + int64_t num_heads,
1058 + int64_t head_size,
1059 + int64_t head_size_v,
1060 + int64_t num_kv_splits,
1061 + int64_t q_strideM,
1062 + int64_t q_strideH,
1063 + int64_t k_strideN,
1064 + int64_t k_strideH,
1065 + int64_t v_strideN,
1066 + int64_t v_strideH,
1067 + float sm_scale,
1068 + float logit_cap,
1069 + int64_t max_num_reqs,
1070 + int64_t max_context_len,
1071 + int64_t max_total_num_tokens,
1072 + int64_t sliding_window_size,
1073 + bool is_cross_attn,
1074 + bool has_encoder_lens,
1075 + bool has_sink) {
1076 + using Vec = at::vec::Vectorized<float>;
1077 +
1078 + // strides
1079 + const int64_t l_stride1 = num_kv_splits * (head_size_v + 1);
1080 + const int64_t l_stride2 = head_size_v + 1;
1081 +
1082 + const bool has_logit_cap = logit_cap > 0;
1083 + float rlogit_cap = has_logit_cap ? 1 / logit_cap : 0.f;
1084 +
1085 + // parallel on [batches, num_heads, num_kv_splits]
1086 + at::parallel_for(0, batches * num_heads * num_kv_splits, 0, [&](int64_t begin, int64_t end) {
1087 + int64_t bs{0}, head_id{0}, kv_id{0};
1088 + data_index_init(begin, bs, batches, head_id, num_heads, kv_id, num_kv_splits);
1089 +
1090 + // s_prime and s_delta
1091 + alignas(64) float s_i[BLOCK_N];
1092 + float* __restrict__ s_delta = s_i;
1093 +
1094 + for (int64_t i = begin; i < end; ++i) {
1095 + // get query
1096 + const scalar_t* __restrict__ q_ptr = query + bs * q_strideM + head_id * q_strideH;
1097 +
1098 + // get key/value
1099 + int64_t seq_len_kv = is_cross_attn ? encoder_lens[bs] : seq_lens[bs];
1100 + int64_t req_pool_id = req_pool_indices[bs];
1101 + int64_t kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0;
1102 + if (sliding_window_size > 0 && seq_len_kv > sliding_window_size) {
1103 + kv_offset = seq_len_kv - sliding_window_size;
1104 + seq_len_kv = sliding_window_size;
1105 + }
1106 + TORCH_CHECK(seq_len_kv <= max_context_len, "seq_len_kv out of scope!");
1107 + TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!");
1108 +
1109 + const int64_t SPLIT_SIZE = div_up(seq_len_kv, num_kv_splits);
1110 + const int64_t kv_start = kv_id * SPLIT_SIZE;
1111 + const int64_t kv_end = std::min(kv_start + SPLIT_SIZE, seq_len_kv);
1112 +
1113 + float m_prime = -std::numeric_limits<float>::infinity();
1114 + float s_prime = 0.f;
1115 +
1116 + // get v_prime, and init to zero
1117 + float* __restrict__ v_prime = attn_logits + i * (head_size_v + 1);
1118 + fill_stub(v_prime, 0.f, head_size_v);
1119 +
1120 + // loop over K and V sequence with BLOCK_N
1121 + for (int64_t n = kv_start; n < kv_end; n += BLOCK_N) {
1122 + int64_t n_size = std::min(BLOCK_N, kv_end - n);
1123 +
1124 + // calculate s_i <- scale * Q @ K
1125 + index_gemm_kernel_nt<scalar_t, index_t>(
1126 + /* A */ q_ptr,
1127 + /* B */ k_buffer + head_id * k_strideH,
1128 + /* C */ s_i,
1129 + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
1130 + /* scl */ sm_scale,
1131 + /* M */ 1,
1132 + /* N */ n_size,
1133 + /* K */ head_size,
1134 + /* lda */ 1,
1135 + /* ldb */ k_strideN,
1136 + /* ldc */ 1,
1137 + /* mtt */ max_total_num_tokens);
1138 +
1139 + // TODO: `tanh` from torch uses sleef u10, going to be slow
1140 + if (has_logit_cap) {
1141 + at::vec::map<float>(
1142 + [logit_cap, rlogit_cap](Vec x) { return Vec(logit_cap) * (x * Vec(rlogit_cap)).tanh(); },
1143 + s_i,
1144 + s_i,
1145 + n_size);
1146 + }
1147 +
1148 + // m_i: max value per row
1149 + float m_i = at::vec::reduce_all<float>([](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, s_i, n_size);
1150 + m_i = std::max(m_i, m_prime);
1151 +
1152 + // m_delta <- exp(m' - m_i)
1153 + float m_delta = std::exp(m_prime - m_i);
1154 +
1155 + // s_delta <- exp(s_i - m_i)
1156 + at::vec::map<float>([m_i](Vec x) { return (x - Vec(m_i)).exp_u20(); }, s_delta, s_i, n_size);
1157 +
1158 + // s' <- s' * m_delta + sum(s_delta)
1159 + s_prime *= m_delta;
1160 + s_prime += at::vec::reduce_all<float>([](Vec& x, Vec& y) { return x + y; }, s_delta, n_size);
1161 +
1162 + m_prime = m_i;
1163 +
1164 + // calculate V' <- s_delta @ V + V' * m_delta
1165 + index_gemm_kernel_nn<scalar_t, index_t>(
1166 + /* A */ s_delta,
1167 + /* B */ v_buffer + head_id * v_strideH,
1168 + /* C */ v_prime,
1169 + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
1170 + /* scl */ &m_delta,
1171 + /* M */ 1,
1172 + /* N */ head_size_v,
1173 + /* K */ n_size,
1174 + /* lda */ 1,
1175 + /* ldb */ v_strideN,
1176 + /* ldc */ 1,
1177 + /* mtt */ max_total_num_tokens);
1178 + } // loop with KV blocks
1179 +
1180 + // only update v' when kv_split_size > 0
1181 + if (kv_end > kv_start) {
1182 + float s = 1 / s_prime;
1183 + at::vec::map<float>([s](Vec out) { return out * Vec(s); }, v_prime, v_prime, head_size_v);
1184 +
1185 + v_prime[head_size_v] = m_prime + std::log(s_prime);
1186 + } else {
1187 + v_prime[head_size_v] = -std::numeric_limits<float>::infinity();
1188 + }
1189 +
1190 + // move to the next index
1191 + data_index_step(bs, batches, head_id, num_heads, kv_id, num_kv_splits);
1192 + }
1193 + });
1194 +
1195 + decode_accumulate_kv_splits(
1196 + output, attn_logits, sinks, batches, num_heads, head_size_v, num_kv_splits, l_stride1, l_stride2, has_sink);
1197 +} // MHA
1198 +
1199 +template <typename scalar_t, typename index_t, int64_t BLOCK_N>
1200 +void decode_attention_mla_kernel_impl(
1201 + scalar_t* __restrict__ output,
1202 + float* __restrict__ attn_logits,
1203 + const scalar_t* __restrict__ query,
1204 + const scalar_t* __restrict__ k_buffer,
1205 + const scalar_t* __restrict__ v_buffer,
1206 + const index_t* __restrict__ req_to_token,
1207 + const int64_t* __restrict__ req_pool_indices,
1208 + const int64_t* __restrict__ seq_lens,
1209 + scalar_t* __restrict__ buffer,
1210 + const scalar_t* __restrict__ sinks,
1211 + int64_t batches,
1212 + int64_t num_heads,
1213 + int64_t head_size,
1214 + int64_t head_size_v,
1215 + int64_t num_kv_splits,
1216 + int64_t q_strideM,
1217 + int64_t q_strideH,
1218 + int64_t k_strideN,
1219 + int64_t k_strideH,
1220 + int64_t v_strideN,
1221 + int64_t v_strideH,
1222 + float sm_scale,
1223 + float logit_cap,
1224 + int64_t max_num_reqs,
1225 + int64_t max_context_len,
1226 + int64_t max_total_num_tokens,
1227 + int64_t buffer_size_per_thread,
1228 + bool has_sink) {
1229 + using Vec = at::vec::Vectorized<float>;
1230 +
1231 + // block length for heads
1232 + const int64_t BLOCK_H = batches == 1 ? 6 : (batches > 16 ? 22 : 11);
1233 +
1234 + // strides
1235 + const int64_t l_stride0 = num_heads * num_kv_splits * (head_size_v + 1);
1236 + const int64_t l_stride1 = num_kv_splits * (head_size_v + 1);
1237 + const int64_t l_stride2 = head_size_v + 1;
1238 +
1239 + TORCH_CHECK(logit_cap == 0.f, "decode MLA: expect no logit_cap.");
1240 +
1241 + // partition the heads into blocks for parallel
1242 + const int64_t num_blocks = div_up(num_heads, BLOCK_H);
1243 +
1244 + // parallel on [batches, num_blocks, num_kv_splits]
1245 + at::parallel_for(0, batches * num_blocks * num_kv_splits, 0, [&](int64_t begin, int64_t end) {
1246 + int64_t bs{0}, block_id{0}, kv_id{0};
1247 + data_index_init(begin, bs, batches, block_id, num_blocks, kv_id, num_kv_splits);
1248 +
1249 + int tid = at::get_thread_num();
1250 + scalar_t* __restrict__ Btmp0 = buffer + tid * buffer_size_per_thread;
1251 + scalar_t* __restrict__ Btmp1 = Btmp0 + BLOCK_N * head_size;
1252 +
1253 + // init Btmp1 just once for each thread to prevent NaN
1254 + // Btmp0 is not needed as it computes full K every single time
1255 + fill_stub(Btmp1, 0.f, BLOCK_N * head_size_v);
1256 +
1257 + alignas(64) float s_i[BLOCK_H * BLOCK_N];
1258 + float* __restrict__ s_delta = s_i;
1259 + alignas(64) scalar_t s_delta2[BLOCK_H * BLOCK_N];
1260 +
1261 + alignas(64) float s_prime[BLOCK_H];
1262 + alignas(64) float m_prime[BLOCK_H];
1263 + alignas(64) float m_delta[BLOCK_H];
1264 +
1265 + for (int64_t i = begin; i < end; ++i) {
1266 + const int64_t h_start = block_id * BLOCK_H;
1267 + const int64_t h_end = std::min(block_id * BLOCK_H + BLOCK_H, num_heads);
1268 + const int64_t h_size = h_end - h_start;
1269 +
1270 + // get query
1271 + const scalar_t* __restrict__ q_ptr = query + bs * q_strideM + h_start * q_strideH;
1272 +
1273 + int64_t seq_len_kv = seq_lens[bs];
1274 + int64_t req_pool_id = req_pool_indices[bs];
1275 + TORCH_CHECK(seq_len_kv <= max_context_len, "seq_len_kv out of scope!");
1276 + TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!");
1277 +
1278 + const int64_t SPLIT_SIZE = div_up(seq_len_kv, num_kv_splits);
1279 + const int64_t kv_start = kv_id * SPLIT_SIZE;
1280 + const int64_t kv_end = std::min(kv_start + SPLIT_SIZE, seq_len_kv);
1281 +
1282 + fill_stub(s_prime, 0.f, BLOCK_H);
1283 + fill_stub(m_prime, -std::numeric_limits<float>::infinity(), BLOCK_H);
1284 +
1285 + // get v_prime, and init to zero
1286 + float* __restrict__ v_prime = attn_logits + bs * l_stride0 + h_start * l_stride1 + kv_id * l_stride2;
1287 + for (int64_t h = 0; h < h_size; ++h) {
1288 + fill_stub(v_prime + h * l_stride1, 0.f, head_size_v);
1289 + }
1290 +
1291 + // loop over K and V sequence with BLOCK_N
1292 + for (int64_t n = kv_start; n < kv_end; n += BLOCK_N) {
1293 + int64_t n_size = std::min(BLOCK_N, kv_end - n);
1294 + const int64_t padded_n_size = div_up(int(n_size), TILE_K) * TILE_K;
1295 +
1296 + // get key and pack
1297 + pack_vnni<scalar_t, index_t>(
1298 + /* dst0 */ Btmp0,
1299 + /* dst1 */ Btmp1,
1300 + /* src */ k_buffer + /* head_kv_id */ 0 * k_strideH,
1301 + /* ind */ req_to_token + req_pool_id * max_context_len + n,
1302 + /* N */ n_size,
1303 + /* K */ head_size,
1304 + /* Kv */ head_size_v,
1305 + /* ld_src */ k_strideN,
1306 + /* ld_dst0 */ BLOCK_N,
1307 + /* ld_dst1 */ head_size_v);
1308 +
1309 + // calculate s_i <- Q @ K
1310 + at::native::cpublas::brgemm(
1311 + /* M */ h_size,
1312 + /* N */ n_size,
1313 + /* K */ head_size,
1314 + /* lda */ q_strideH,
1315 + /* ldb */ BLOCK_N,
1316 + /* ldc */ BLOCK_N,
1317 + /* add_C */ false,
1318 + /* A */ q_ptr,
1319 + /* B */ Btmp0,
1320 + /* C */ s_i);
1321 +
1322 + const Vec scale_vec = Vec(sm_scale);
1323 + for (int64_t h = 0; h < h_size; ++h) {
1324 + // s_i <- s_i * scale
1325 + at::vec::map<float>(
1326 + [scale_vec](Vec x) { return x * scale_vec; }, s_i + h * BLOCK_N, s_i + h * BLOCK_N, n_size);
1327 +
1328 + // m_i: max value per row
1329 + float m_i = at::vec::reduce_all<float>(
1330 + [](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, s_i + h * BLOCK_N, n_size);
1331 + m_i = std::max(m_i, m_prime[h]);
1332 +
1333 + // m_delta <- exp(m' - m_i)
1334 + m_delta[h] = std::exp(m_prime[h] - m_i);
1335 +
1336 + // s_delta <- exp(s_i - m_i)
1337 + at::vec::map<float>(
1338 + [m_i](Vec x) { return (x - Vec(m_i)).exp_u20(); }, s_delta + h * BLOCK_N, s_i + h * BLOCK_N, n_size);
1339 +
1340 + // s' <- s' * m_delta + sum(s_delta)
1341 + s_prime[h] *= m_delta[h];
1342 + s_prime[h] += at::vec::reduce_all<float>([](Vec& x, Vec& y) { return x + y; }, s_delta + h * BLOCK_N, n_size);
1343 +
1344 + m_prime[h] = m_i;
1345 +
1346 + // v' <- v' * m_delta
1347 + float scale_m = m_delta[h];
1348 + at::vec::map<float>(
1349 + [scale_m](Vec x) { return x * Vec(scale_m); },
1350 + v_prime + h * l_stride1,
1351 + v_prime + h * l_stride1,
1352 + head_size_v);
1353 +
1354 + // pad s_delta with 0 first and then convert to scalar_t
1355 + fill_stub(s_delta + h * BLOCK_N + n_size, 0.f, padded_n_size - n_size);
1356 + copy_stub<scalar_t, BLOCK_N>(s_delta2 + h * BLOCK_N, s_delta + h * BLOCK_N);
1357 + }
1358 +
1359 + // calculate V' <- s_delta @ V + V'
1360 + at::native::cpublas::brgemm(
1361 + /* M */ h_size,
1362 + /* N */ head_size_v,
1363 + /* K */ padded_n_size, // n_size
1364 + /* lda */ BLOCK_N,
1365 + /* ldb */ head_size_v,
1366 + /* ldc */ l_stride1,
1367 + /* add_C */ true,
1368 + /* A */ s_delta2,
1369 + /* B */ Btmp1,
1370 + /* C */ v_prime);
1371 + } // loop with KV blocks
1372 +
1373 + // only update v' when kv_split_size > 0
1374 + if (kv_end > kv_start) {
1375 + for (int64_t h = 0; h < h_size; ++h) {
1376 + float s = 1 / s_prime[h];
1377 + at::vec::map<float>(
1378 + [s](Vec out) { return out * Vec(s); }, v_prime + h * l_stride1, v_prime + h * l_stride1, head_size_v);
1379 + (v_prime + h * l_stride1)[head_size_v] = m_prime[h] + std::log(s_prime[h]);
1380 + }
1381 + } else {
1382 + for (int64_t h = 0; h < h_size; ++h) {
1383 + (v_prime + h * l_stride1)[head_size_v] = -std::numeric_limits<float>::infinity();
1384 + }
1385 + }
1386 +
1387 + // move to the next index
1388 + data_index_step(bs, batches, block_id, num_blocks, kv_id, num_kv_splits);
1389 + }
1390 + at::native::cpublas::brgemm_release();
1391 + });
1392 +
1393 + decode_accumulate_kv_splits(
1394 + output, attn_logits, sinks, batches, num_heads, head_size_v, num_kv_splits, l_stride1, l_stride2, has_sink);
1395 +} // MLA
1396 +
1397 +template <typename scalar_t, typename index_t, int64_t BLOCK_N>
1398 +void decode_attention_grouped_kernel_impl(
1399 + scalar_t* __restrict__ output,
1400 + float* __restrict__ attn_logits,
1401 + const scalar_t* __restrict__ query,
1402 + const scalar_t* __restrict__ k_buffer,
1403 + const scalar_t* __restrict__ v_buffer,
1404 + const index_t* __restrict__ req_to_token,
1405 + const int64_t* __restrict__ req_pool_indices,
1406 + const int64_t* __restrict__ seq_lens,
1407 + const int64_t* __restrict__ encoder_lens,
1408 + const scalar_t* __restrict__ sinks,
1409 + int64_t batches,
1410 + int64_t num_heads,
1411 + int64_t num_heads_kv,
1412 + int64_t head_size,
1413 + int64_t head_size_v,
1414 + int64_t num_kv_splits,
1415 + int64_t q_strideM,
1416 + int64_t q_strideH,
1417 + int64_t k_strideN,
1418 + int64_t k_strideH,
1419 + int64_t v_strideN,
1420 + int64_t v_strideH,
1421 + float sm_scale,
1422 + float logit_cap,
1423 + int64_t max_num_reqs,
1424 + int64_t max_context_len,
1425 + int64_t max_total_num_tokens,
1426 + int64_t sliding_window_size,
1427 + bool is_cross_attn,
1428 + bool has_encoder_lens,
1429 + bool has_sink) {
1430 + using Vec = at::vec::Vectorized<float>;
1431 +
1432 + // block length for heads
1433 + // we parallel on [batches, divup(num_heads, BLOCK_H), num_kv_splits]
1434 + // use smaller BLOCK_H when batches is small to utilize all cores
1435 + constexpr int64_t kBLOCK_H = 16;
1436 + const int64_t BLOCK_H = std::min(4 * batches, kBLOCK_H);
1437 +
1438 + // strides
1439 + const int64_t l_stride0 = num_heads * num_kv_splits * (head_size_v + 1);
1440 + const int64_t l_stride1 = num_kv_splits * (head_size_v + 1);
1441 + const int64_t l_stride2 = head_size_v + 1;
1442 +
1443 + const bool has_logit_cap = logit_cap > 0;
1444 + float rlogit_cap = has_logit_cap ? 1 / logit_cap : 0.f;
1445 +
1446 + // partition the heads into blocks for parallel
1447 + const int64_t num_groups = num_heads / num_heads_kv;
1448 + const int64_t num_blocks = div_up(num_groups, BLOCK_H);
1449 +
1450 + // parallel on [batches, num_heads_kv, num_blocks, num_kv_splits]
1451 + at::parallel_for(0, batches * num_heads_kv * num_blocks * num_kv_splits, 0, [&](int64_t begin, int64_t end) {
1452 + int64_t bs{0}, head_kv_id{0}, block_id{0}, kv_id{0};
1453 + data_index_init(begin, bs, batches, head_kv_id, num_heads_kv, block_id, num_blocks, kv_id, num_kv_splits);
1454 +
1455 + alignas(64) float s_i[BLOCK_H * BLOCK_N];
1456 + float* __restrict__ s_delta = s_i;
1457 +
1458 + alignas(64) float s_prime[BLOCK_H];
1459 + alignas(64) float m_prime[BLOCK_H];
1460 + alignas(64) float m_delta[BLOCK_H];
1461 +
1462 + for (int64_t i = begin; i < end; ++i) {
1463 + const int64_t h_start = head_kv_id * num_groups + block_id * BLOCK_H;
1464 + const int64_t h_end = head_kv_id * num_groups + std::min(block_id * BLOCK_H + BLOCK_H, num_groups);
1465 + const int64_t h_size = h_end - h_start;
1466 +
1467 + // get query
1468 + const scalar_t* __restrict__ q_ptr = query + bs * q_strideM + h_start * q_strideH;
1469 +
1470 + int64_t seq_len_kv = is_cross_attn ? encoder_lens[bs] : seq_lens[bs];
1471 + int64_t req_pool_id = req_pool_indices[bs];
1472 + int64_t kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0;
1473 + TORCH_CHECK(seq_len_kv <= max_context_len, "seq_len_kv out of scope!");
1474 + TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!");
1475 + if (sliding_window_size > 0 && seq_len_kv > sliding_window_size) {
1476 + kv_offset = seq_len_kv - sliding_window_size;
1477 + seq_len_kv = sliding_window_size;
1478 + }
1479 + const int64_t SPLIT_SIZE = div_up(seq_len_kv, num_kv_splits);
1480 + const int64_t kv_start = kv_id * SPLIT_SIZE;
1481 + const int64_t kv_end = std::min(kv_start + SPLIT_SIZE, seq_len_kv);
1482 +
1483 + fill_stub(s_prime, 0.f, BLOCK_H);
1484 + fill_stub(m_prime, -std::numeric_limits<float>::infinity(), BLOCK_H);
1485 +
1486 + // get v_prime, and init to zero
1487 + float* __restrict__ v_prime = attn_logits + bs * l_stride0 + h_start * l_stride1 + kv_id * l_stride2;
1488 + for (int64_t h = 0; h < h_size; ++h) {
1489 + fill_stub(v_prime + h * l_stride1, 0.f, head_size_v);
1490 + }
1491 +
1492 + // loop over K and V sequence with BLOCK_N
1493 + for (int64_t n = kv_start; n < kv_end; n += BLOCK_N) {
1494 + int64_t n_size = std::min(BLOCK_N, kv_end - n);
1495 +
1496 + // calculate Q @ K
1497 + index_gemm_kernel_nt<scalar_t, index_t>(
1498 + /* A */ q_ptr,
1499 + /* B */ k_buffer + head_kv_id * k_strideH,
1500 + /* C */ s_i,
1501 + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
1502 + /* scl */ sm_scale,
1503 + /* M */ h_size,
1504 + /* N */ n_size,
1505 + /* K */ head_size,
1506 + /* lda */ q_strideH,
1507 + /* ldb */ k_strideN,
1508 + /* ldc */ BLOCK_N,
1509 + /* mtt */ max_total_num_tokens);
1510 +
1511 + if (has_logit_cap) {
1512 + at::vec::map<float>(
1513 + [logit_cap, rlogit_cap](Vec x) { return Vec(logit_cap) * (x * Vec(rlogit_cap)).tanh(); },
1514 + s_i,
1515 + s_i,
1516 + BLOCK_H * BLOCK_N);
1517 + }
1518 +
1519 + // update the sm_scale coefficients
1520 + for (int64_t h = 0; h < h_size; ++h) {
1521 + // m_i: max value per row
1522 + float m_i = at::vec::reduce_all<float>(
1523 + [](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, s_i + h * BLOCK_N, n_size);
1524 + m_i = std::max(m_i, m_prime[h]);
1525 +
1526 + // m_delta <- exp(m' - m_i)
1527 + m_delta[h] = std::exp(m_prime[h] - m_i);
1528 +
1529 + // s_delta <- exp(s_i - m_i)
1530 + at::vec::map<float>(
1531 + [m_i](Vec x) { return (x - Vec(m_i)).exp_u20(); }, s_delta + h * BLOCK_N, s_i + h * BLOCK_N, n_size);
1532 +
1533 + // s' <- s' * m_delta + sum(s_delta)
1534 + s_prime[h] *= m_delta[h];
1535 + s_prime[h] += at::vec::reduce_all<float>([](Vec& x, Vec& y) { return x + y; }, s_delta + h * BLOCK_N, n_size);
1536 +
1537 + m_prime[h] = m_i;
1538 + }
1539 +
1540 + // calculate V' <- s_delta @ V + V' * m_delta
1541 + index_gemm_kernel_nn<scalar_t, index_t>(
1542 + /* A */ s_delta,
1543 + /* B */ v_buffer + head_kv_id * v_strideH,
1544 + /* C */ v_prime,
1545 + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
1546 + /* scl */ m_delta,
1547 + /* M */ h_size,
1548 + /* N */ head_size_v,
1549 + /* K */ n_size,
1550 + /* lda */ BLOCK_N,
1551 + /* ldb */ v_strideN,
1552 + /* ldc */ l_stride1,
1553 + /* mtt */ max_total_num_tokens);
1554 + } // loop with KV blocks
1555 +
1556 + // only update v' when kv_split_size > 0
1557 + if (kv_end > kv_start) {
1558 + for (int64_t h = 0; h < h_size; ++h) {
1559 + float s = 1 / s_prime[h];
1560 + at::vec::map<float>(
1561 + [s](Vec out) { return out * Vec(s); }, v_prime + h * l_stride1, v_prime + h * l_stride1, head_size_v);
1562 + (v_prime + h * l_stride1)[head_size_v] = m_prime[h] + std::log(s_prime[h]);
1563 + }
1564 + } else {
1565 + for (int64_t h = 0; h < h_size; ++h) {
1566 + (v_prime + h * l_stride1)[head_size_v] = -std::numeric_limits<float>::infinity();
1567 + }
1568 + }
1569 +
1570 + // move to the next index
1571 + data_index_step(bs, batches, head_kv_id, num_heads_kv, block_id, num_blocks, kv_id, num_kv_splits);
1572 + }
1573 + });
1574 +
1575 + decode_accumulate_kv_splits(
1576 + output, attn_logits, sinks, batches, num_heads, head_size_v, num_kv_splits, l_stride1, l_stride2, has_sink);
1577 +} // GQA/MQA
1578 +
1579 +} // anonymous namespace
1580 +
1581 +// query: [num_tokens, num_heads, head_size]
1582 +// output: [num_tokens, num_heads, head_size]
1583 +// k_buffer: [max_total_num_tokens, num_heads, head_size]
1584 +// v_buffer: [max_total_num_tokens, num_heads, head_size_v]
1585 +// attn_logits: [num_seqs, num_heads, num_kv_splits, head_size_v + 1]
1586 +// req_to_token: [max_num_reqs, max_context_len] int32 or int64
1587 +// req_pool_indices: [num_seqs] int64
1588 +// seq_lens: [num_seqs] int64
1589 +// encoder_lens: [num_seqs] int64 or None
1590 +// sinks: [num_heads] or None
1591 +void decode_attention_cpu(
1592 + at::Tensor& query,
1593 + at::Tensor& k_buffer,
1594 + at::Tensor& v_buffer,
1595 + at::Tensor& output,
1596 + const std::optional<at::Tensor>& key,
1597 + const std::optional<at::Tensor>& value,
1598 + const std::optional<at::Tensor>& loc,
1599 + at::Tensor& attn_logits,
1600 + at::Tensor& req_to_token,
1601 + at::Tensor& req_pool_indices,
1602 + at::Tensor& seq_lens,
1603 + double sm_scale,
1604 + double logit_cap,
1605 + bool is_cross_attn,
1606 + int64_t sliding_window_size,
1607 + std::optional<at::Tensor> encoder_lens,
1608 + std::optional<at::Tensor> sinks) {
1609 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(query);
1610 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_buffer);
1611 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_buffer);
1612 + CHECK_DIM(3, query);
1613 + CHECK_DIM(3, k_buffer);
1614 + CHECK_DIM(3, v_buffer);
1615 +
1616 + int64_t num_seqs = seq_lens.size(0);
1617 + int64_t max_num_reqs = req_to_token.size(0);
1618 + int64_t max_context_len = req_to_token.size(1);
1619 + int64_t max_total_num_tokens = k_buffer.size(0);
1620 +
1621 + int64_t num_heads = query.size(1);
1622 + int64_t num_heads_kv = k_buffer.size(1);
1623 + int64_t head_size = query.size(2);
1624 + int64_t head_size_v = v_buffer.size(2);
1625 +
1626 + int64_t num_kv_splits = attn_logits.size(2);
1627 +
1628 + CHECK_EQ(attn_logits.size(0), num_seqs);
1629 + CHECK_EQ(attn_logits.size(1), num_heads);
1630 + CHECK_EQ(attn_logits.size(3), head_size_v + 1);
1631 + CHECK_EQ(attn_logits.scalar_type(), at::kFloat);
1632 +
1633 + // strides for query
1634 + int64_t q_strideM = query.stride(0);
1635 + int64_t q_strideH = query.stride(1);
1636 +
1637 + // strides for k_buffer and v_buffer
1638 + int64_t k_strideN = k_buffer.stride(0);
1639 + int64_t k_strideH = k_buffer.stride(1);
1640 + int64_t v_strideN = v_buffer.stride(0);
1641 + int64_t v_strideH = v_buffer.stride(1);
1642 +
1643 + // check index data types
1644 + const auto index_dtype = req_to_token.scalar_type();
1645 + TORCH_CHECK(
1646 + index_dtype == at::kInt || index_dtype == at::kLong,
1647 + "decode: expect req_to_token to be int32 or int64, got ",
1648 + index_dtype);
1649 + TORCH_CHECK(seq_lens.scalar_type() == at::kLong, "decode: expect req_lens to be int64, got ", seq_lens.scalar_type());
1650 + TORCH_CHECK(
1651 + req_pool_indices.scalar_type() == at::kLong,
1652 + "decode: expect req_pool_indices to be int64, got ",
1653 + req_pool_indices.scalar_type());
1654 +
1655 + // check if we have MLA here
1656 + void* k_buffer_data = k_buffer.data_ptr();
1657 + void* v_buffer_data = v_buffer.data_ptr();
1658 + const bool is_mla = (k_buffer_data == v_buffer_data) && (num_heads_kv == 1) && (head_size == head_size_v + 64);
1659 +
1660 + // block length for k_buffer and v_buffer
1661 + constexpr int BLOCK_N = 256;
1662 +
1663 + // buffer for packing k_cache and v_cache
1664 + int num_threads = at::get_num_threads();
1665 + int64_t size_per_thread = is_mla ? BLOCK_N * head_size + BLOCK_N * head_size_v : 0;
1666 + auto buffer = at::empty({num_threads, size_per_thread}, k_buffer.options());
1667 + bool has_encoder_lens = encoder_lens.has_value();
1668 + // Since encoder_lens is not used when it is None, encoder_lens_t can be initialized as any tensor of int64_t dtype.
1669 + at::Tensor encoder_lens_t = seq_lens;
1670 + if (has_encoder_lens) {
1671 + encoder_lens_t = encoder_lens.value();
1672 + CHECK_EQ(encoder_lens_t.size(0), num_seqs);
1673 + }
1674 + bool has_sink = sinks.has_value();
1675 + at::Tensor sinks_tensor = has_sink ? sinks.value() : at::empty({num_heads}, query.options());
1676 + CHECK_DIM(1, sinks_tensor);
1677 + CHECK_EQ(sinks_tensor.size(0), num_heads);
1678 + AT_DISPATCH_REDUCED_FLOATING_TYPES(query.scalar_type(), "decode_attention_kernel", [&] {
1679 + AT_DISPATCH_INDEX_TYPES(index_dtype, "decode_attention_indices", [&] {
1680 + if (key.has_value()) {
1681 + TORCH_CHECK(value.has_value(), "key and value should have values at the same time")
1682 + TORCH_CHECK(loc.has_value(), "loc must be given when key/value are given")
1683 + auto loc_tensor = loc.value();
1684 + CHECK_DIM(1, loc_tensor);
1685 + CHECK_EQ(loc_tensor.numel(), num_seqs);
1686 + auto key_tensor = key.value();
1687 + auto value_tensor = value.value();
1688 + // for MLA, key and value shares the same storage and value could be non-contiguous
1689 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(key_tensor);
1690 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(value_tensor);
1691 + CHECK_DIM(3, key_tensor);
1692 + CHECK_DIM(3, value_tensor);
1693 + // strides for new key and value
1694 + int64_t nk_strideN = key_tensor.stride(0);
1695 + int64_t nk_strideH = key_tensor.stride(1);
1696 + int64_t nv_strideN = value_tensor.stride(0);
1697 + int64_t nv_strideH = value_tensor.stride(1);
1698 + // update the kv buffer
1699 + decode_set_kv_buffer(
1700 + (scalar_t*)k_buffer_data,
1701 + (scalar_t*)v_buffer_data,
1702 + key_tensor.data_ptr<scalar_t>(),
1703 + value_tensor.data_ptr<scalar_t>(),
1704 + loc_tensor.data_ptr<int64_t>(),
1705 + num_seqs,
1706 + num_heads_kv,
1707 + head_size,
1708 + head_size_v,
1709 + k_strideN,
1710 + k_strideH,
1711 + v_strideN,
1712 + v_strideH,
1713 + nk_strideN,
1714 + nk_strideH,
1715 + nv_strideN,
1716 + nv_strideH,
1717 + is_mla);
1718 + }
1719 +
1720 + if (num_heads == num_heads_kv) {
1721 + // MHA
1722 + decode_attention_kernel_impl<scalar_t, index_t, BLOCK_N>(
1723 + output.data_ptr<scalar_t>(),
1724 + attn_logits.data_ptr<float>(),
1725 + query.data_ptr<scalar_t>(),
1726 + (const scalar_t*)k_buffer_data,
1727 + (const scalar_t*)v_buffer_data,
1728 + req_to_token.data_ptr<index_t>(),
1729 + req_pool_indices.data_ptr<int64_t>(),
1730 + seq_lens.data_ptr<int64_t>(),
1731 + encoder_lens_t.data_ptr<int64_t>(),
1732 + sinks_tensor.data_ptr<scalar_t>(),
1733 + num_seqs,
1734 + num_heads,
1735 + head_size,
1736 + head_size_v,
1737 + num_kv_splits,
1738 + q_strideM,
1739 + q_strideH,
1740 + k_strideN,
1741 + k_strideH,
1742 + v_strideN,
1743 + v_strideH,
1744 + sm_scale,
1745 + logit_cap,
1746 + max_num_reqs,
1747 + max_context_len,
1748 + max_total_num_tokens,
1749 + sliding_window_size,
1750 + is_cross_attn,
1751 + has_encoder_lens,
1752 + has_sink);
1753 + } else if (is_mla) {
1754 + // MLA
1755 + decode_attention_mla_kernel_impl<scalar_t, index_t, BLOCK_N>(
1756 + output.data_ptr<scalar_t>(),
1757 + attn_logits.data_ptr<float>(),
1758 + query.data_ptr<scalar_t>(),
1759 + (const scalar_t*)k_buffer_data,
1760 + (const scalar_t*)v_buffer_data,
1761 + req_to_token.data_ptr<index_t>(),
1762 + req_pool_indices.data_ptr<int64_t>(),
1763 + seq_lens.data_ptr<int64_t>(),
1764 + buffer.data_ptr<scalar_t>(),
1765 + sinks_tensor.data_ptr<scalar_t>(),
1766 + num_seqs,
1767 + num_heads,
1768 + head_size,
1769 + head_size_v,
1770 + num_kv_splits,
1771 + q_strideM,
1772 + q_strideH,
1773 + k_strideN,
1774 + k_strideH,
1775 + v_strideN,
1776 + v_strideH,
1777 + sm_scale,
1778 + logit_cap,
1779 + max_num_reqs,
1780 + max_context_len,
1781 + max_total_num_tokens,
1782 + size_per_thread,
1783 + has_sink);
1784 + } else {
1785 + // GQA/MQA
1786 + decode_attention_grouped_kernel_impl<scalar_t, index_t, BLOCK_N>(
1787 + output.data_ptr<scalar_t>(),
1788 + attn_logits.data_ptr<float>(),
1789 + query.data_ptr<scalar_t>(),
1790 + (const scalar_t*)k_buffer_data,
1791 + (const scalar_t*)v_buffer_data,
1792 + req_to_token.data_ptr<index_t>(),
1793 + req_pool_indices.data_ptr<int64_t>(),
1794 + seq_lens.data_ptr<int64_t>(),
1795 + encoder_lens_t.data_ptr<int64_t>(),
1796 + sinks_tensor.data_ptr<scalar_t>(),
1797 + num_seqs,
1798 + num_heads,
1799 + num_heads_kv,
1800 + head_size,
1801 + head_size_v,
1802 + num_kv_splits,
1803 + q_strideM,
1804 + q_strideH,
1805 + k_strideN,
1806 + k_strideH,
1807 + v_strideN,
1808 + v_strideH,
1809 + sm_scale,
1810 + logit_cap,
1811 + max_num_reqs,
1812 + max_context_len,
1813 + max_total_num_tokens,
1814 + sliding_window_size,
1815 + is_cross_attn,
1816 + has_encoder_lens,
1817 + has_sink);
1818 + }
1819 + });
1820 + });
1821 +}
csrc/cpu/sgl-kernels/extend.cpp new
+568
@@ -0,0 +1,568 @@
1 +// Adapted from
2 +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3 +
4 +// clang-format off
5 +
6 +#include "common.h"
7 +#include "flash_attn.h"
8 +#include "gemm.h"
9 +
10 +namespace {
11 +
12 +// [NOTE]: extend attention for CPU
13 +// 1. BLOCK_M and BLOCK_N tuned for various seq lengths
14 +// 2. can handle non-contiguous k_extend and v_extend
15 +// 3. computes attention for prefix and extend separately
16 +// 4. TODO: apply head dimension blocking to optimize GQA
17 +// 5. optional tree mask for speculative decoding TARGET_VERIFY (EAGLE topk > 1):
18 +// `tree_mask` is a flat [batches * qlen * qlen] bool tensor in
19 +// TreeMaskMode::QLEN_ONLY layout, where qlen == extend_seq_lens[bs] ==
20 +// max_len_extend (uniform across the batch, equal to draft_token_num).
21 +// Row i = query draft token, column j = key draft token; true means query i
22 +// may attend key j (each row marks self + ancestors + root). The committed
23 +// prefix (stage 1) is implicitly fully visible to every draft token, which
24 +// is why the mask only covers the qlen x qlen new-token block; the GPU
25 +// FULL_MASK layout carries the prefix columns explicitly but they are
26 +// all-true for EAGLE. When tree_mask is absent, stage 2 falls back to the
27 +// plain causal mask (correct for non-spec extend and topk == 1 chains).
28 +//
29 +
30 +template <typename scalar_t, typename index_t, int BLOCK_M, int BLOCK_N>
31 +void extend_attention_kernel_impl(
32 + scalar_t* __restrict__ o_extend,
33 + const scalar_t* __restrict__ q_extend,
34 + const scalar_t* __restrict__ k_extend,
35 + const scalar_t* __restrict__ v_extend,
36 + const scalar_t* __restrict__ k_buffer,
37 + const scalar_t* __restrict__ v_buffer,
38 + const index_t* __restrict__ req_to_token,
39 + const int64_t* __restrict__ req_pool_indices,
40 + const int64_t* __restrict__ seq_lens,
41 + const int64_t* __restrict__ encoder_lens,
42 + const index_t* __restrict__ extend_seq_lens,
43 + const index_t* __restrict__ extend_start_loc,
44 + const void* __restrict__ buffer,
45 + const scalar_t* __restrict__ sinks,
46 + const bool* __restrict__ tree_mask,
47 + int batches,
48 + int num_heads,
49 + int num_heads_kv,
50 + int head_size,
51 + int head_size_v,
52 + int q_strideM,
53 + int q_strideH,
54 + int ke_strideN,
55 + int ke_strideH,
56 + int ve_strideN,
57 + int ve_strideH,
58 + int k_strideN,
59 + int k_strideH,
60 + int v_strideN,
61 + int v_strideH,
62 + float sm_scale,
63 + int max_num_reqs,
64 + int max_context_len,
65 + int max_total_num_tokens,
66 + int max_len_extend,
67 + int buffer_size_per_thread,
68 + int64_t sliding_window_size,
69 + bool is_prefix_skipped,
70 + bool is_cross_attn,
71 + bool has_encoder_lens,
72 + bool has_sink) {
73 + // strides
74 + const int o_strideM = num_heads * head_size_v;
75 + const int o_strideH = head_size_v;
76 +
77 + // we use same buffer for packed key and value
78 + const int ldb_tmp = std::max(head_size, head_size_v);
79 +
80 + const int num_groups = num_heads / num_heads_kv;
81 + TORCH_CHECK(num_groups * num_heads_kv == num_heads);
82 +
83 + // number of blocks along M
84 + int MB = div_up(max_len_extend, BLOCK_M);
85 +
86 + // parallel on [batches, num_heads, BM]
87 + at::parallel_for(0, batches * num_heads * MB, 0, [&](int begin, int end) {
88 + int bs{0}, head_id{0}, mb{0};
89 + data_index_init(begin, bs, batches, head_id, num_heads, mb, MB);
90 +
91 + int tid = at::get_thread_num();
92 + // s_i: [BLOCK_M, BLOCK_N]
93 + float* __restrict__ s_i = reinterpret_cast<float*>((char*)(buffer) + tid * buffer_size_per_thread);
94 +
95 + // v_prime: [BLOCK_M, head_size_v]
96 + float* __restrict__ v_prime = s_i + BLOCK_M * BLOCK_N;
97 +
98 + // s_delta: [BLOCK_M, BLOCK_N]
99 + scalar_t* __restrict__ s_delta = reinterpret_cast<scalar_t*>(v_prime + BLOCK_M * head_size_v);
100 +
101 + // Btmp: [BLOCK_N, max(head_size, head_size_v)]
102 + scalar_t* __restrict__ Btmp = reinterpret_cast<scalar_t*>(s_delta + BLOCK_M * BLOCK_N);
103 +
104 + // init Btmp just once for each thread to prevent NaN
105 + fill_stub(Btmp, 0.f, BLOCK_N * ldb_tmp);
106 + fill_stub(s_delta, 0.f, BLOCK_M * BLOCK_N);
107 +
108 + alignas(64) float s_prime[BLOCK_M];
109 + alignas(64) float m_prime[BLOCK_M];
110 +
111 + for (int i = begin; i < end; ++i) {
112 + // seq_len = prefix + extend
113 + int head_kv_id = head_id / num_groups;
114 + int seq_len = seq_lens[bs];
115 + int seq_len_extend = extend_seq_lens[bs];
116 + int seq_len_prefix = seq_len - seq_len_extend;
117 + int seq_extend_start_loc = extend_start_loc[bs];
118 +
119 + int req_pool_id = req_pool_indices[bs];
120 + int kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0;
121 + TORCH_CHECK(seq_len_prefix >= 0, "prefix len < 0!");
122 + TORCH_CHECK(seq_len <= max_context_len, "seq_len out of scope!");
123 + TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!");
124 +
125 + if (is_prefix_skipped) {
126 + TORCH_CHECK(seq_len_prefix == 0, "extend attention: expect seq_len_prefix to be 0, got ", seq_len_prefix);
127 + }
128 +
129 + if (tree_mask != nullptr) {
130 + // QLEN_ONLY layout assumes a uniform qlen across the batch (TARGET_VERIFY)
131 + TORCH_CHECK(
132 + seq_len_extend == max_len_extend,
133 + "extend attention: tree_mask requires uniform extend_seq_lens, got ",
134 + seq_len_extend,
135 + " vs ",
136 + max_len_extend);
137 + }
138 +
139 + // offset and size in MB
140 + int m = mb * BLOCK_M;
141 + int m_size = std::min(BLOCK_M, seq_len_extend - m);
142 +
143 + if (m_size <= 0) {
144 + data_index_step(bs, batches, head_id, num_heads, mb, MB);
145 + continue;
146 + }
147 +
148 + // get query
149 + const scalar_t* __restrict__ q_ptr = q_extend + (seq_extend_start_loc + m) * q_strideM + head_id * q_strideH;
150 +
151 + // init v', s' and m'
152 + fill_stub(v_prime, 0.f, m_size * head_size_v);
153 + fill_stub(s_prime, 0.f, m_size);
154 + fill_stub(m_prime, -std::numeric_limits<scalar_t>::infinity(), m_size);
155 + // stage 1: compute scores with prefix
156 + int kv_start = 0;
157 + int kv_end = is_cross_attn ? encoder_lens[bs] : seq_len_prefix;
158 + for (int n = kv_start; n < kv_end; n += BLOCK_N) {
159 + int n_size = std::min(BLOCK_N, kv_end - n);
160 +
161 + // `n_size` is K in 2nd gemm, pad to TILE_K;
162 + const int padded_n_size = div_up(n_size, TILE_K) * TILE_K;
163 +
164 + // get key and pack
165 + pack_vnni<scalar_t, index_t>(
166 + /* dst */ Btmp,
167 + /* src */ k_buffer + head_kv_id * k_strideH,
168 + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
169 + /* N */ n_size,
170 + /* K */ head_size,
171 + /* ld_src */ k_strideN,
172 + /* ld_dst */ BLOCK_N);
173 +
174 + // calculate s_i <- Q @ K
175 + at::native::cpublas::brgemm(
176 + /* M */ m_size,
177 + /* N */ n_size,
178 + /* K */ head_size,
179 + /* lda */ q_strideM,
180 + /* ldb */ BLOCK_N,
181 + /* ldc */ BLOCK_N,
182 + /* add_C */ false,
183 + /* A */ q_ptr,
184 + /* B */ Btmp,
185 + /* C */ s_i);
186 +
187 + for (int row = 0; row < m_size; ++row) {
188 + if (sliding_window_size > 0) {
189 + int last_col = seq_len_prefix + row + m - sliding_window_size + 1;
190 + if (last_col >= n + n_size) {
191 + continue;
192 + }
193 + fill_stub(s_i + row * BLOCK_N, -std::numeric_limits<float>::infinity(), last_col - n);
194 + }
195 + flash_attn_softmax<scalar_t, BLOCK_M, BLOCK_N>::apply(
196 + s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale, row);
197 + }
198 +
199 + // get value and pack
200 + pack_vnni2<scalar_t>(
201 + /* dst */ Btmp,
202 + /* src */ v_buffer + head_kv_id * v_strideH,
203 + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
204 + /* K */ n_size,
205 + /* N */ head_size_v,
206 + /* ld_src */ v_strideN,
207 + /* ld_dst */ head_size_v);
208 +
209 + // calculate V' <- s_delta @ V + V'
210 + at::native::cpublas::brgemm(
211 + /* M */ m_size,
212 + /* N */ head_size_v,
213 + /* K */ padded_n_size, // n_size
214 + /* lda */ BLOCK_N,
215 + /* ldb */ head_size_v,
216 + /* ldc */ head_size_v,
217 + /* add_C */ true,
218 + /* A */ s_delta,
219 + /* B */ Btmp,
220 + /* C */ v_prime);
221 + } // loop with seq_len_prefix
222 + if (!is_cross_attn) {
223 + // stage 2: compute the triangle part
224 + int num_keys = std::min(seq_len_extend, m + BLOCK_M);
225 + for (int n = 0; n < num_keys; n += BLOCK_N) {
226 + int n_size = std::min(BLOCK_N, num_keys - n);
227 +
228 + // `n_size` is K in 2nd gemm, pad to TILE_K;
229 + const int padded_n_size = div_up(n_size, TILE_K) * TILE_K;
230 +
231 + // get key and pack
232 + pack_vnni<scalar_t>(
233 + /* dst */ Btmp,
234 + /* src */ k_extend + (seq_extend_start_loc + n) * ke_strideN + head_kv_id * ke_strideH,
235 + /* N */ n_size,
236 + /* K */ head_size,
237 + /* ld_src */ ke_strideN,
238 + /* ld_dst */ BLOCK_N);
239 +
240 + // calculate s_i <- Q @ K
241 + at::native::cpublas::brgemm(
242 + /* M */ m_size,
243 + /* N */ n_size,
244 + /* K */ head_size,
245 + /* lda */ q_strideM,
246 + /* ldb */ BLOCK_N,
247 + /* ldc */ BLOCK_N,
248 + /* add_C */ false,
249 + /* A */ q_ptr,
250 + /* B */ Btmp,
251 + /* C */ s_i);
252 +
253 + // apply tree mask (speculative TARGET_VERIFY) or causal mask
254 + if (tree_mask != nullptr) {
255 + // [Note] tree mask for EAGLE topk > 1 (TreeMaskMode::QLEN_ONLY).
256 + // mask[bs][m + row][n + col] == false -> query draft token (m + row)
257 + // may not attend key draft token (n + col); set the score to -inf
258 + // before softmax. The tree mask subsumes the causal constraint:
259 + // ancestors always precede descendants in the draft token ordering,
260 + // so permitted keys satisfy j <= i and the causal `num_keys` bound
261 + // above remains valid.
262 + const bool* __restrict__ mask_base =
263 + tree_mask + (static_cast<int64_t>(bs) * seq_len_extend + m) * seq_len_extend + n;
264 + for (int row = 0; row < m_size; ++row) {
265 + float* __restrict__ row_ptr = s_i + row * BLOCK_N;
266 + const bool* __restrict__ mask_ptr = mask_base + static_cast<int64_t>(row) * seq_len_extend;
267 + for (int col = 0; col < n_size; ++col) {
268 + if (!mask_ptr[col]) {
269 + row_ptr[col] = -std::numeric_limits<float>::infinity();
270 + }
271 + }
272 + }
273 + } else if (n + n_size - 1 > m) {
274 + // apply causal mask
275 + // [Note] condition to apply causal mask.
276 + // Mask any block whose last key (n + n_size - 1) is strictly after the first query position (m), i.e. n +
277 + // n_size - 1 > m. The original condition was `num_keys - n <= BLOCK_N` (last n-block only). That was
278 + // correct when BLOCK_M <= BLOCK_N/2 because earlier n-blocks were guaranteed to contain only past keys.
279 + // With BLOCK_M=512, BLOCK_N=768:
280 + // BLOCK_M > BLOCK_N/2, so the first n-block can contain future keys.
281 + // Example: m=512 (mb=1), num_keys=1024, first n-block covers keys [0, 768).
282 + // Query row=0 is at position 512, so keys 513..767 are future and must be
283 + // masked — but `num_keys - 0 = 1024 > BLOCK_N` skips masking entirely,
284 + // producing wrong (non-causal) attention for rows 0..254 of this m-block.
285 + for (int row = 0; row < m_size; ++row) {
286 + int last_col = m + row - n;
287 + // [Note] mask the entire row if last_col < 0.
288 + // Clamp to -1: when n > m + row every key in this block is a future
289 + // key, so the entire row should be masked. Without this clamp,
290 + // last_col+1 <= 0 and fill_stub would write before row_ptr.
291 + last_col = std::max(last_col, -1);
292 + // fill [last_col + 1, n_size) to -inf
293 + float* row_ptr = s_i + row * BLOCK_N;
294 + fill_stub(row_ptr + last_col + 1, -std::numeric_limits<float>::infinity(), n_size - last_col - 1);
295 + }
296 + }
297 +
298 + for (int row = 0; row < m_size; ++row) {
299 + if (sliding_window_size > 0 && row + m + 1 >= n + sliding_window_size - 1 &&
300 + row + m + 1 < n + sliding_window_size + n_size) {
301 + fill_stub(
302 + s_i + row * BLOCK_N, -std::numeric_limits<float>::infinity(), row + m - n - sliding_window_size + 1);
303 + } else if (sliding_window_size > 0 && row + m + 1 >= n + sliding_window_size) {
304 + continue;
305 + }
306 + flash_attn_softmax<scalar_t, BLOCK_M, BLOCK_N>::apply(
307 + s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale, row);
308 + }
309 +
310 + // get value and pack
311 + pack_vnni2<scalar_t>(
312 + /* dst */ Btmp,
313 + /* src */ v_extend + (seq_extend_start_loc + n) * ve_strideN + head_kv_id * ve_strideH,
314 + /* K */ n_size,
315 + /* N */ head_size_v,
316 + /* ld_src */ ve_strideN,
317 + /* ld_dst */ head_size_v);
318 +
319 + // calculate V' <- s_delta @ V + V'
320 + at::native::cpublas::brgemm(
321 + /* M */ m_size,
322 + /* N */ head_size_v,
323 + /* K */ padded_n_size, // n_size
324 + /* lda */ BLOCK_N,
325 + /* ldb */ head_size_v,
326 + /* ldc */ head_size_v,
327 + /* add_C */ true,
328 + /* A */ s_delta,
329 + /* B */ Btmp,
330 + /* C */ v_prime);
331 + } // loop with seq_len_extend
332 + }
333 + scalar_t* __restrict__ out_ptr = o_extend + (seq_extend_start_loc + m) * o_strideM + head_id * o_strideH;
334 + for (int row = 0; row < m_size; ++row) {
335 + if (has_sink) {
336 + s_prime[row] += std::exp(sinks[head_id] - m_prime[row]);
337 + }
338 + float s = 1 / s_prime[row];
339 + copy_stub<scalar_t>(out_ptr + row * o_strideM, v_prime + row * head_size_v, s, head_size_v);
340 + }
341 +
342 + // move to the next index
343 + data_index_step(bs, batches, head_id, num_heads, mb, MB);
344 + }
345 + at::native::cpublas::brgemm_release();
346 + });
347 +}
348 +
349 +} // anonymous namespace
350 +
351 +template <int BLOCK_M, int BLOCK_N>
352 +inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int head_size_v) {
353 + static_assert(BLOCK_M <= BLOCK_N, "Make sure BLOCK_M <= BLOCK_N to prevent buffer overflows during causal masking");
354 + const int size_per_thread =
355 + /* s_i */ BLOCK_M * BLOCK_N * sizeof(float) +
356 + /* v_prime */ BLOCK_M * head_size_v * sizeof(float) +
357 + /* s_delta */ BLOCK_M * BLOCK_N * sizeof(uint16_t) +
358 + /* Btmp */ BLOCK_N * std::max(head_size, head_size_v) * sizeof(uint16_t);
359 +
360 + buffer.resize_({num_threads, size_per_thread});
361 + return size_per_thread;
362 +}
363 +
364 +#define LAUNCH_EXTEND_ATTENTION_KERNEL(BLOCK_M, BLOCK_N) \
365 + do { \
366 + int sz = resize_buffer<BLOCK_M, BLOCK_N>(buffer, num_threads, head_size, head_size_v); \
367 + \
368 + extend_attention_kernel_impl<scalar_t, index_t, BLOCK_M, BLOCK_N>( \
369 + o_extend.data_ptr<scalar_t>(), \
370 + q_extend.data_ptr<scalar_t>(), \
371 + k_extend.data_ptr<scalar_t>(), \
372 + v_extend.data_ptr<scalar_t>(), \
373 + k_buffer.data_ptr<scalar_t>(), \
374 + v_buffer.data_ptr<scalar_t>(), \
375 + req_to_token.data_ptr<index_t>(), \
376 + req_pool_indices.data_ptr<int64_t>(), \
377 + seq_lens.data_ptr<int64_t>(), \
378 + encoder_lens_t.data_ptr<int64_t>(), \
379 + extend_seq_lens.data_ptr<index_t>(), \
380 + extend_start_loc.data_ptr<index_t>(), \
381 + buffer.data_ptr(), \
382 + sinks_tensor.data_ptr<scalar_t>(), \
383 + tree_mask_ptr, \
384 + num_seqs, \
385 + num_heads, \
386 + num_heads_kv, \
387 + head_size, \
388 + head_size_v, \
389 + q_strideM, \
390 + q_strideH, \
391 + ke_strideN, \
392 + ke_strideH, \
393 + ve_strideN, \
394 + ve_strideH, \
395 + k_strideN, \
396 + k_strideH, \
397 + v_strideN, \
398 + v_strideH, \
399 + sm_scale, \
400 + max_num_reqs, \
401 + max_context_len, \
402 + max_total_num_tokens, \
403 + max_len_extend, \
404 + sz, \
405 + sliding_window_size, \
406 + is_prefix_skipped, \
407 + is_cross_attn, \
408 + has_encoder_lens, \
409 + has_sink); \
410 + } while (0)
411 +
412 +// q_extend, k_extend, v_extend, o_extend: contiguous tensors
413 +// k_buffer, v_buffer: (prefix + extend) tensors in mem_manager
414 +//
415 +// q_extend: [num_tokens, num_heads, head_size]
416 +// k_extend: [num_extend_tokens, num_heads, head_size]
417 +// v_extend: [num_extend_tokens, num_heads, head_size]
418 +// o_extend: [num_tokens, num_heads, head_size]
419 +// k_buffer: [max_total_num_tokens, num_heads, head_size]
420 +// v_buffer: [max_total_num_tokens, num_heads, head_size]
421 +// req_to_token: [max_num_reqs, max_context_len] int32 or int64
422 +// req_pool_indices: [num_seqs] int64
423 +// seq_lens: [num_seqs] int64
424 +// extend_seq_lens: [num_seqs]
425 +// extend_start_loc: [num_seqs]
426 +// encoder_lens: [num_seqs] int64 or None
427 +// sinks: [num_heads] or None
428 +// tree_mask: [num_seqs * max_len_extend * max_len_extend] bool or None
429 +// TreeMaskMode::QLEN_ONLY tree mask for speculative TARGET_VERIFY; see [NOTE] 5 above.
430 +void extend_attention_cpu(
431 + at::Tensor& q_extend,
432 + const std::optional<at::Tensor>& k_extend_opt,
433 + const std::optional<at::Tensor>& v_extend_opt,
434 + at::Tensor& o_extend,
435 + at::Tensor& k_buffer,
436 + at::Tensor& v_buffer,
437 + at::Tensor& req_to_token,
438 + at::Tensor& req_pool_indices,
439 + at::Tensor& seq_lens,
440 + at::Tensor& extend_seq_lens,
441 + at::Tensor& extend_start_loc,
442 + int64_t max_len_extend,
443 + double sm_scale,
444 + double logit_cap,
445 + bool is_cross_attn,
446 + int64_t sliding_window_size,
447 + std::optional<at::Tensor> encoder_lens,
448 + std::optional<at::Tensor> sinks,
449 + std::optional<at::Tensor> tree_mask) {
450 + if (!is_cross_attn) {
451 + TORCH_CHECK(
452 + k_extend_opt.has_value() && v_extend_opt.has_value(),
453 + "k_extend and v_extend are required for non-cross attention");
454 + }
455 + // Since k_extend and v_extend are not used for cross attention, they can be initialized as k_buffer and v_buffer
456 + // here.
457 + auto k_extend = k_extend_opt.has_value() ? k_extend_opt.value() : k_buffer;
458 + auto v_extend = v_extend_opt.has_value() ? v_extend_opt.value() : v_buffer;
459 +
460 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(q_extend);
461 + CHECK_INPUT(o_extend);
462 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_extend);
463 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_extend);
464 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_buffer);
465 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_buffer);
466 +
467 + int num_seqs = seq_lens.size(0);
468 + int max_num_reqs = req_to_token.size(0);
469 + int max_context_len = req_to_token.size(1);
470 + int max_total_num_tokens = k_buffer.size(0);
471 +
472 + int num_heads = q_extend.size(1);
473 + int num_heads_kv = k_extend.size(1);
474 + int head_size = q_extend.size(2);
475 + int head_size_v = v_extend.size(2);
476 +
477 + // strides for q_extend, k_extend and v_extend
478 + int q_strideM = q_extend.stride(0);
479 + int q_strideH = q_extend.stride(1);
480 + int ke_strideN = k_extend.stride(0);
481 + int ke_strideH = k_extend.stride(1);
482 + int ve_strideN = v_extend.stride(0);
483 + int ve_strideH = v_extend.stride(1);
484 +
485 + // strides for k_buffer and v_buffer
486 + int k_strideN = k_buffer.stride(0);
487 + int k_strideH = k_buffer.stride(1);
488 + int v_strideN = v_buffer.stride(0);
489 + int v_strideH = v_buffer.stride(1);
490 +
491 + // check sizes
492 + CHECK_EQ(req_pool_indices.size(0), num_seqs);
493 + CHECK_EQ(extend_seq_lens.size(0), num_seqs);
494 + CHECK_EQ(extend_start_loc.size(0), num_seqs);
495 + CHECK_EQ(v_extend.size(1), num_heads_kv);
496 + CHECK_EQ(k_buffer.size(1), v_buffer.size(1));
497 +
498 + // MLA will skip prefix part
499 + const bool is_prefix_skipped = k_buffer.size(1) != num_heads_kv;
500 +
501 + // check index data types
502 + const auto index_dtype = req_to_token.scalar_type();
503 + TORCH_CHECK(
504 + index_dtype == at::kInt || index_dtype == at::kLong,
505 + "extend: expect req_to_token to be int32 or int64, got ",
506 + index_dtype);
507 + TORCH_CHECK(seq_lens.scalar_type() == at::kLong, "extend: expect req_lens to be int64, got ", seq_lens.scalar_type());
508 + TORCH_CHECK(
509 + req_pool_indices.scalar_type() == at::kLong,
510 + "extend: expect req_pool_indices to be int64, got ",
511 + req_pool_indices.scalar_type());
512 + TORCH_CHECK(
513 + extend_seq_lens.scalar_type() == index_dtype && extend_start_loc.scalar_type() == index_dtype,
514 + "extend: expect extend_seq_lens and extend_start_loc to have same dtype as req_to_token.");
515 +
516 + // D and DV need to be 32x as we transpose by 512-bit
517 + TORCH_CHECK(head_size % 32 == 0, "invalid head_size ", head_size);
518 + TORCH_CHECK(head_size_v % 32 == 0, "invalid head_size_v ", head_size_v);
519 +
520 + int num_threads = at::get_num_threads();
521 + auto buffer = at::empty({}, q_extend.options().dtype(at::kChar));
522 +
523 + bool has_encoder_lens = encoder_lens.has_value();
524 + // Since encoder_lens is not used when it is None, encoder_lens_t can be initialized as any tensor of int64_t dtype.
525 + at::Tensor encoder_lens_t = seq_lens;
526 + if (has_encoder_lens) {
527 + encoder_lens_t = encoder_lens.value();
528 + CHECK_EQ(encoder_lens_t.size(0), num_seqs);
529 + }
530 + bool has_sink = sinks.has_value();
531 + at::Tensor sinks_tensor = has_sink ? sinks.value() : at::empty({num_heads}, q_extend.options());
532 + CHECK_DIM(1, sinks_tensor);
533 + CHECK_EQ(sinks_tensor.size(0), num_heads);
534 +
535 + const bool* tree_mask_ptr = nullptr;
536 + if (tree_mask.has_value()) {
537 + const at::Tensor& tree_mask_t = tree_mask.value();
538 + CHECK_INPUT(tree_mask_t);
539 + TORCH_CHECK(
540 + tree_mask_t.scalar_type() == at::kBool, "extend: expect tree_mask to be bool, got ", tree_mask_t.scalar_type());
541 + TORCH_CHECK(
542 + tree_mask_t.numel() == static_cast<int64_t>(num_seqs) * max_len_extend * max_len_extend,
543 + "extend: expect tree_mask numel to be num_seqs * max_len_extend^2 = ",
544 + static_cast<int64_t>(num_seqs) * max_len_extend * max_len_extend,
545 + ", got ",
546 + tree_mask_t.numel());
547 + TORCH_CHECK(!is_cross_attn, "extend: tree_mask is not supported for cross attention");
548 + // The window mask derives query positions from the row index
549 + // (seq_len_prefix + m + row), but tree-mask rows sit at their tree depth,
550 + // which is <= the row index; combining the two would over-mask the prefix.
551 + TORCH_CHECK(sliding_window_size <= 0, "extend: tree_mask is not supported with sliding window attention");
552 + tree_mask_ptr = tree_mask_t.data_ptr<bool>();
553 + }
554 +
555 + AT_DISPATCH_REDUCED_FLOATING_TYPES(q_extend.scalar_type(), "extend_attention_kernel", [&] {
556 + AT_DISPATCH_INDEX_TYPES(index_dtype, "extend_attention_indices", [&] {
557 + if (max_len_extend <= 256) {
558 + LAUNCH_EXTEND_ATTENTION_KERNEL(32, 64);
559 + } else if (max_len_extend <= 1024) {
560 + LAUNCH_EXTEND_ATTENTION_KERNEL(128, 256);
561 + } else if (max_len_extend <= 4096) {
562 + LAUNCH_EXTEND_ATTENTION_KERNEL(256, 768);
563 + } else { // max_len_extend > 4096
564 + LAUNCH_EXTEND_ATTENTION_KERNEL(512, 768);
565 + }
566 + });
567 + });
568 +}
csrc/cpu/sgl-kernels/flash_attn.h new
+250
@@ -0,0 +1,250 @@
1 +// Adapted from
2 +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3 +
4 +// clang-format off
5 +
6 +#pragma once
7 +#include "common.h"
8 +#include "vec.h"
9 +#include "vec_pack.h"
10 +
11 +template <typename scalar_t>
12 +inline void fill_stub(scalar_t* __restrict__ out, float val, int size) {
13 + using Vec = at::vec::Vectorized<scalar_t>;
14 + constexpr int kVecSize = Vec::size();
15 + const Vec data_vec = Vec(static_cast<scalar_t>(val));
16 + int d = 0;
17 +#pragma GCC unroll 4
18 + for (; d <= size - kVecSize; d += kVecSize) {
19 + data_vec.store(out + d);
20 + }
21 + if (size - d > 0) {
22 + data_vec.store(out + d, size - d);
23 + }
24 +}
25 +
26 +template <typename scalar_t, int BLOCK_N>
27 +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input) {
28 + static_assert(BLOCK_N % 32 == 0);
29 + using bVec = at::vec::Vectorized<scalar_t>;
30 + using fVec = at::vec::Vectorized<float>;
31 +
32 + constexpr int COLS = BLOCK_N / 16;
33 + auto store = [&](auto i) {
34 + constexpr int col = i % COLS;
35 + // for COLS = 2, 4 use 512bit store
36 + if constexpr (col % 2 == 0) {
37 + auto [a_fvec0, a_fvec1] = load_float_vec2(input + col * 16);
38 + bVec out_bvec = convert_from_float_ext<scalar_t>(a_fvec0, a_fvec1);
39 + out_bvec.store(out + col * 16);
40 + }
41 + };
42 + Unroll<COLS>{}(store);
43 +}
44 +
45 +template <typename scalar_t>
46 +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ acc, float s, int size) {
47 + using bVec = at::vec::Vectorized<scalar_t>;
48 + using fVec = at::vec::Vectorized<float>;
49 + constexpr int kVecSize = bVec::size();
50 + const fVec s_fvec = fVec(s);
51 + int d = 0;
52 +#pragma GCC unroll 4
53 + for (; d <= size - kVecSize; d += kVecSize) {
54 + auto [a_fvec0, a_fvec1] = load_float_vec2(acc + d);
55 + a_fvec0 = a_fvec0 * s_fvec;
56 + a_fvec1 = a_fvec1 * s_fvec;
57 + bVec out_bvec = convert_from_float_ext<scalar_t>(a_fvec0, a_fvec1);
58 + out_bvec.store(out + d);
59 + }
60 + for (; d < size; ++d) {
61 + out[d] = static_cast<scalar_t>(acc[d] * s);
62 + }
63 +}
64 +
65 +#if defined(CPU_CAPABILITY_AVX512)
66 +template <>
67 +inline void copy_stub<at::BFloat16>(at::BFloat16* __restrict__ out, const float* __restrict__ acc, float s, int size) {
68 + const __m512 vscale = _mm512_set1_ps(s);
69 + int d = 0;
70 +#pragma GCC unroll 4
71 + for (; d <= size - 32; d += 32) {
72 + __m512 va0 = _mm512_mul_ps(_mm512_loadu_ps(acc + d), vscale);
73 + __m512 va1 = _mm512_mul_ps(_mm512_loadu_ps(acc + d + 16), vscale);
74 + __m512i vb = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0));
75 + _mm512_storeu_si512(out + d, vb);
76 + }
77 + int remainder = size - d;
78 + if (remainder > 0) {
79 + if (remainder <= 16) {
80 + const __mmask16 vmask = (1ULL << remainder) - 1;
81 + __m512 va = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask, acc + d), vscale);
82 + __m256i vb = (__m256i)(_mm512_cvtneps_pbh(va));
83 + _mm256_mask_storeu_epi16(reinterpret_cast<__m256i*>(out + d), vmask, vb);
84 + } else { // remainder > 16
85 + const __mmask16 vmask = (1ULL << (remainder - 16)) - 1;
86 + __m512 va0 = _mm512_mul_ps(_mm512_loadu_ps(acc + d), vscale);
87 + __m512 va1 = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask, acc + d + 16), vscale);
88 + __m512i vb = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0));
89 + const __mmask32 vmask2 = (1ULL << remainder) - 1;
90 + _mm512_mask_storeu_epi16(reinterpret_cast<__m512i*>(out + d), vmask2, vb);
91 + }
92 + }
93 +}
94 +#endif
95 +
96 +template <typename scalar_t, int BLOCK_M, int BLOCK_N>
97 +struct flash_attn_softmax {
98 + static inline void apply(
99 + float* __restrict__ s_i,
100 + scalar_t* __restrict__ s_delta2,
101 + float* __restrict__ v_prime,
102 + float* __restrict__ s_prime,
103 + float* __restrict__ m_prime,
104 + int m_size,
105 + int n_size,
106 + int padded_n_size,
107 + int head_size_v,
108 + const float sm_scale,
109 + int row) {
110 + using Vec = at::vec::Vectorized<float>;
111 + const Vec scale_vec = Vec(sm_scale);
112 + float* s_delta = s_i;
113 + // s_i <- s_i * scale
114 + at::vec::map<float>([scale_vec](Vec x) { return x * scale_vec; }, s_i + row * BLOCK_N, s_i + row * BLOCK_N, n_size);
115 +
116 + // m_i: max value per row
117 + float m_i =
118 + at::vec::reduce_all<float>([](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, s_i + row * BLOCK_N, n_size);
119 + m_i = std::max(m_i, m_prime[row]);
120 +
121 + // m_delta <- exp(m' - m_i)
122 + float m_delta = std::exp(m_prime[row] - m_i);
123 +
124 + // s_delta <- exp(s_i - m_i)
125 + at::vec::map<float>(
126 + [m_i](Vec x) { return (x - Vec(m_i)).fexp_u20(); }, s_delta + row * BLOCK_N, s_i + row * BLOCK_N, n_size);
127 +
128 + // s' <- s' * m_delta + sum(s_delta)
129 + s_prime[row] *= m_delta;
130 + s_prime[row] += at::vec::reduce_all<float>([](Vec& x, Vec& y) { return x + y; }, s_delta + row * BLOCK_N, n_size);
131 +
132 + m_prime[row] = m_i;
133 +
134 + // v' <- v' * m_delta
135 + at::vec::map<float>(
136 + [m_delta](Vec x) { return x * Vec(m_delta); },
137 + v_prime + row * head_size_v,
138 + v_prime + row * head_size_v,
139 + head_size_v);
140 +
141 + // Keep s_delta row-major for the following brgemm(P @ V), and only
142 + // convert the columns that brgemm will consume.
143 + fill_stub(s_delta + row * BLOCK_N + n_size, 0.f, padded_n_size - n_size);
144 + copy_stub<scalar_t>(s_delta2 + row * BLOCK_N, s_delta + row * BLOCK_N, 1.f, padded_n_size);
145 + }
146 +};
147 +
148 +#if defined(CPU_CAPABILITY_AVX512)
149 +template <int BLOCK_M, int BLOCK_N>
150 +struct flash_attn_softmax<at::BFloat16, BLOCK_M, BLOCK_N> {
151 + static inline void apply(
152 + float* __restrict__ s_i,
153 + at::BFloat16* __restrict__ s_delta2,
154 + float* __restrict__ v_prime,
155 + float* __restrict__ s_prime,
156 + float* __restrict__ m_prime,
157 + int m_size,
158 + int n_size,
159 + int padded_n_size,
160 + int head_size_v,
161 + const float sm_scale,
162 + int row) {
163 + float* s_delta = s_i;
164 + const __m512 vscale = _mm512_set1_ps(sm_scale);
165 +
166 + int n_remainder = n_size & 15; // 0xF
167 + const __mmask16 vmask = (1ULL << n_remainder) - 1;
168 +
169 + int v_remainder = head_size_v & 15; // 0xF
170 + const __mmask16 vmask1 = (1ULL << v_remainder) - 1;
171 +
172 + constexpr float NEG_INF = -std::numeric_limits<float>::infinity();
173 +
174 + __m512 va;
175 + __m256i vb;
176 + __m512 vmax;
177 + __m512 vsum;
178 + __m512 vmdelta;
179 +
180 + const __m512 vneg_inf = _mm512_set1_ps(NEG_INF);
181 +
182 + int m = row;
183 + vmax = vneg_inf;
184 +
185 + // s_i <- s_i * scale
186 + int n = 0;
187 + for (; n <= n_size - 16; n += 16) {
188 + va = _mm512_mul_ps(_mm512_loadu_ps(s_i + m * BLOCK_N + n), vscale);
189 + vmax = _mm512_max_ps(va, vmax);
190 + }
191 + if (n_remainder > 0) {
192 + va = _mm512_mul_ps(_mm512_mask_loadu_ps(vneg_inf, vmask, s_i + m * BLOCK_N + n), vscale);
193 + vmax = _mm512_max_ps(va, vmax);
194 + }
195 +
196 + // m_i: max value per row
197 + float m_i = _mm512_reduce_max_ps(vmax);
198 + m_i = std::max(m_i, m_prime[m]);
199 + vmax = _mm512_set1_ps(m_i);
200 +
201 + // m_delta <- exp(m' - m_i)
202 + float m_delta = std::exp(m_prime[m] - m_i);
203 +
204 + // s_delta <- exp(s_i - m_i)
205 + vsum = _mm512_setzero_ps();
206 + for (n = 0; n <= n_size - 16; n += 16) {
207 + va = _mm512_mul_ps(_mm512_loadu_ps(s_i + m * BLOCK_N + n), vscale);
208 + va = _mm512_fexp_u20_ps(_mm512_sub_ps(va, vmax));
209 + vsum = _mm512_add_ps(vsum, va);
210 +
211 + vb = (__m256i)(_mm512_cvtneps_pbh(va));
212 + _mm256_storeu_si256(reinterpret_cast<__m256i*>(s_delta2 + m * BLOCK_N + n), vb);
213 + }
214 + if (n_remainder > 0) {
215 + va = _mm512_mul_ps(_mm512_mask_loadu_ps(vneg_inf, vmask, s_i + m * BLOCK_N + n), vscale);
216 + va = _mm512_fexp_u20_ps(_mm512_sub_ps(va, vmax));
217 + vsum = _mm512_add_ps(vsum, va);
218 +
219 + vb = (__m256i)(_mm512_cvtneps_pbh(va));
220 + _mm256_mask_storeu_epi16(reinterpret_cast<__m256i*>(s_delta2 + m * BLOCK_N + n), vmask, vb);
221 + }
222 +
223 + // s' <- s' * m_delta + sum(s_delta)
224 + s_prime[m] *= m_delta;
225 + s_prime[m] += _mm512_reduce_add_ps(vsum);
226 +
227 + m_prime[m] = m_i;
228 +
229 + // pad s_delta with 0, pad_size range from [0, 32)
230 + int pad_size = padded_n_size - n_size;
231 + if (pad_size > 0) {
232 + const __m512i vzero = _mm512_setzero_si512();
233 + __mmask32 vmask2 = (1ULL << pad_size) - 1;
234 + _mm512_mask_storeu_epi16(reinterpret_cast<__m512i*>(s_delta2 + m * BLOCK_N + n_size), vmask2, vzero);
235 + }
236 +
237 + // v' <- v' * m_delta
238 + vmdelta = _mm512_set1_ps(m_delta);
239 + int k = 0;
240 + for (; k <= head_size_v - 16; k += 16) {
241 + va = _mm512_mul_ps(_mm512_loadu_ps(v_prime + m * head_size_v + k), vmdelta);
242 + _mm512_storeu_ps(reinterpret_cast<__m512*>(v_prime + m * head_size_v + k), va);
243 + }
244 + if (v_remainder > 0) {
245 + va = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask1, v_prime + m * head_size_v + k), vmdelta);
246 + _mm512_mask_storeu_ps(reinterpret_cast<__m512*>(v_prime + m * head_size_v + k), vmask1, va);
247 + }
248 + }
249 +};
250 +#endif
csrc/cpu/sgl-kernels/mla_cache.cpp new
+111
@@ -0,0 +1,111 @@
1 +// vLLM-native CPU cache-write op for MLA's single-latent-buffer KV cache.
2 +//
3 +// `concat_and_cache_mla` (the generic MLA cache-write op used by every GPU
4 +// backend) is registered CUDA-only. This is the CPU counterpart, adapted in
5 +// spirit from SGLang's `store_cache_cpu` (csrc/cpu/kvcache.cpp) but
6 +// generalized to write two source tensors (`kv_c_normed`, `k_pe`) into two
7 +// different column-offset ranges of the SAME destination row -- SGLang's
8 +// version assumes k/v land in two independent, equal-row-width cache
9 +// tensors, which doesn't hold here since MLA's cache is one 576-wide buffer
10 +// and the two column ranges (512-wide, 64-wide) don't match the buffer's
11 +// true per-token stride, so the write can't reuse `store_cache_cpu` as-is.
12 +
13 +#include "common.h"
14 +#include "vec.h"
15 +
16 +namespace {
17 +
18 +template <typename scalar_t>
19 +inline void copy_stub(scalar_t* __restrict__ dst,
20 + const scalar_t* __restrict__ src, int64_t size) {
21 + int64_t d = 0;
22 +#if defined(CPU_CAPABILITY_AVX512)
23 + using Vec = at::vec::Vectorized<scalar_t>;
24 + constexpr int64_t kVecSize = Vec::size();
25 + for (; d <= size - kVecSize; d += kVecSize) {
26 + Vec data = Vec::loadu(src + d);
27 + data.store(dst + d);
28 + }
29 +#endif
30 + for (; d < size; ++d) {
31 + dst[d] = src[d];
32 + }
33 +}
34 +
35 +template <typename scalar_t, typename index_t>
36 +void concat_and_cache_mla_kernel_impl(
37 + const scalar_t* __restrict__ kv_c_normed, // [num_tokens, kv_lora_rank]
38 + const scalar_t* __restrict__ k_pe, // [num_tokens, qk_rope_head_dim]
39 + scalar_t* __restrict__ kv_cache, // [.., kv_lora_rank + qk_rope_head_dim]
40 + const index_t* __restrict__ slot_mapping, // [num_tokens]
41 + int64_t num_tokens, int64_t kv_lora_rank, int64_t qk_rope_head_dim,
42 + int64_t kv_c_stride, int64_t k_pe_stride, int64_t cache_stride) {
43 + at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) {
44 + for (int64_t i = begin; i < end; ++i) {
45 + int64_t slot = static_cast<int64_t>(slot_mapping[i]);
46 + if (slot < 0) {
47 + // padded/invalid token, matches `concat_and_cache_mla`'s semantics.
48 + continue;
49 + }
50 + scalar_t* __restrict__ cache_row = kv_cache + slot * cache_stride;
51 + copy_stub(cache_row, kv_c_normed + i * kv_c_stride, kv_lora_rank);
52 + copy_stub(cache_row + kv_lora_rank, k_pe + i * k_pe_stride,
53 + qk_rope_head_dim);
54 + }
55 + });
56 +}
57 +
58 +} // namespace
59 +
60 +// kv_c_normed : [num_tokens, kv_lora_rank]
61 +// k_pe : [num_tokens, qk_rope_head_dim] or [num_tokens, 1,
62 +// qk_rope_head_dim] kv_cache : [num_blocks, block_size, kv_lora_rank +
63 +// qk_rope_head_dim] slot_mapping: [num_tokens] int32/int64, absolute physical
64 +// row index
65 +// (block_id * block_size + block_offset); negative entries are
66 +// skipped (padded tokens), matching `concat_and_cache_mla`.
67 +void concat_and_cache_mla_cpu(const at::Tensor& kv_c_normed,
68 + const at::Tensor& k_pe, at::Tensor& kv_cache,
69 + const at::Tensor& slot_mapping) {
70 + TORCH_CHECK(kv_c_normed.dim() == 2,
71 + "kv_c_normed must be 2D [num_tokens, kv_lora_rank]");
72 + TORCH_CHECK(k_pe.dim() == 2 || k_pe.dim() == 3, "k_pe must be 2D or 3D");
73 + TORCH_CHECK(kv_cache.dim() == 3,
74 + "kv_cache must be 3D [num_blocks, block_size, head_size]");
75 + TORCH_CHECK(kv_c_normed.stride(-1) == 1,
76 + "kv_c_normed innermost dim must be contiguous");
77 + TORCH_CHECK(k_pe.stride(-1) == 1, "k_pe innermost dim must be contiguous");
78 + TORCH_CHECK(kv_cache.stride(-1) == 1,
79 + "kv_cache innermost dim must be contiguous");
80 +
81 + int64_t num_tokens = kv_c_normed.size(0);
82 + int64_t kv_lora_rank = kv_c_normed.size(1);
83 + int64_t qk_rope_head_dim = k_pe.size(-1);
84 + int64_t head_size = kv_cache.size(-1);
85 + TORCH_CHECK(head_size == kv_lora_rank + qk_rope_head_dim,
86 + "kv_cache head_size must equal kv_lora_rank + qk_rope_head_dim");
87 + TORCH_CHECK(slot_mapping.size(0) == num_tokens, "slot_mapping size mismatch");
88 +
89 + // Real physical per-token stride of the cache, read from the tensor's own
90 + // strides rather than assumed to equal head_size (paged/pooled buffers are
91 + // not guaranteed to have zero inter-row padding in general, even though in
92 + // practice a freshly-allocated MLA cache is fully contiguous).
93 + int64_t cache_stride = kv_cache.stride(1);
94 +
95 + const auto dtype = kv_cache.scalar_type();
96 + TORCH_CHECK(dtype == kv_c_normed.scalar_type() && dtype == k_pe.scalar_type(),
97 + "concat_and_cache_mla_cpu: dtype mismatch");
98 + const auto index_dtype = slot_mapping.scalar_type();
99 + TORCH_CHECK(index_dtype == at::kLong || index_dtype == at::kInt,
100 + "slot_mapping must be int32 or int64");
101 +
102 + AT_DISPATCH_REDUCED_FLOATING_TYPES(dtype, "concat_and_cache_mla_cpu", [&] {
103 + AT_DISPATCH_INDEX_TYPES(index_dtype, "concat_and_cache_mla_cpu_index", [&] {
104 + concat_and_cache_mla_kernel_impl<scalar_t, index_t>(
105 + kv_c_normed.data_ptr<scalar_t>(), k_pe.data_ptr<scalar_t>(),
106 + kv_cache.data_ptr<scalar_t>(), slot_mapping.data_ptr<index_t>(),
107 + num_tokens, kv_lora_rank, qk_rope_head_dim, kv_c_normed.stride(0),
108 + k_pe.stride(0), cache_stride);
109 + });
110 + });
111 +}
csrc/cpu/torch_bindings.cpp
+62
@@ -91,6 +91,39 @@ at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2,
91 const std::optional<at::Tensor>& bias,
92 at::ScalarType out_dtype, bool is_vnni);
93
94 +// Adapted from sglang: MLA CPU kernels (AMX-only)
95 +void decode_attention_cpu(at::Tensor& query, at::Tensor& k_buffer,
96 + at::Tensor& v_buffer, at::Tensor& output,
97 + const std::optional<at::Tensor>& key,
98 + const std::optional<at::Tensor>& value,
99 + const std::optional<at::Tensor>& loc,
100 + at::Tensor& attn_logits, at::Tensor& req_to_token,
101 + at::Tensor& req_pool_indices, at::Tensor& seq_lens,
102 + double sm_scale, double logit_cap, bool is_cross_attn,
103 + int64_t sliding_window_size,
104 + std::optional<at::Tensor> encoder_lens,
105 + std::optional<at::Tensor> sinks);
106 +
107 +void extend_attention_cpu(
108 + at::Tensor& q_extend, const std::optional<at::Tensor>& k_extend,
109 + const std::optional<at::Tensor>& v_extend, at::Tensor& o_extend,
110 + at::Tensor& k_buffer, at::Tensor& v_buffer, at::Tensor& req_to_token,
111 + at::Tensor& req_pool_indices, at::Tensor& seq_lens,
112 + at::Tensor& extend_seq_lens, at::Tensor& extend_start_loc,
113 + int64_t max_len_extend, double sm_scale, double logit_cap,
114 + bool is_cross_attn, int64_t sliding_window_size,
115 + std::optional<at::Tensor> encoder_lens, std::optional<at::Tensor> sinks,
116 + std::optional<at::Tensor> tree_mask);
117 +
118 +void bmm_cpu(at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni,
119 + const std::optional<at::Tensor>& scale);
120 +
121 +// vLLM-native: CPU cache-write op for MLA's single-latent-buffer KV cache
122 +// (the CUDA-only `concat_and_cache_mla` has no CPU dispatch).
123 +void concat_and_cache_mla_cpu(const at::Tensor& kv_c_normed,
124 + const at::Tensor& k_pe, at::Tensor& kv_cache,
125 + const at::Tensor& slot_mapping);
126 +
127 // Adapted from sglang: INT4 W4A8 kernels
128 std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
129 at::Tensor qweight, // awq: (*, K, N / 8) || gptq: (*, K / 8, N) , int32
@@ -521,6 +554,35 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
554 "pad_slot_id, "
555 "bool is_vnni) -> Tensor");
556 ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
557 +
558 + // Adapted from sglang: MLA CPU kernels (AMX-only, DeepSeek V2/V3/R1)
559 + ops.def(
560 + "decode_attention_cpu(Tensor query, Tensor k_buffer, Tensor v_buffer, "
561 + "Tensor(a!) output, Tensor? key, Tensor? value, Tensor? loc, Tensor "
562 + "attn_logits, Tensor req_to_token, Tensor req_pool_indices, Tensor "
563 + "seq_lens, float sm_scale, float logit_cap, bool is_cross_attn, int "
564 + "sliding_window_size, Tensor? encoder_lens, Tensor? sinks) -> ()");
565 + ops.impl("decode_attention_cpu", torch::kCPU, &decode_attention_cpu);
566 +
567 + ops.def(
568 + "extend_attention_cpu(Tensor q_extend, Tensor? k_extend, Tensor? "
569 + "v_extend, Tensor(a!) o_extend, Tensor k_buffer, Tensor v_buffer, "
570 + "Tensor req_to_token, Tensor req_pool_indices, Tensor seq_lens, Tensor "
571 + "extend_seq_lens, Tensor extend_start_loc, int max_len_extend, float "
572 + "sm_scale, float logit_cap, bool is_cross_attn, int "
573 + "sliding_window_size, Tensor? encoder_lens, Tensor? sinks, Tensor? "
574 + "tree_mask=None) -> ()");
575 + ops.impl("extend_attention_cpu", torch::kCPU, &extend_attention_cpu);
576 +
577 + ops.def(
578 + "bmm_cpu(Tensor(a!) out, Tensor mat1, Tensor mat2, bool is_vnni, "
579 + "Tensor? scale) -> ()");
580 + ops.impl("bmm_cpu", torch::kCPU, &bmm_cpu);
581 +
582 + ops.def(
583 + "concat_and_cache_mla_cpu(Tensor kv_c_normed, Tensor k_pe, "
584 + "Tensor(a!) kv_cache, Tensor slot_mapping) -> ()");
585 + ops.impl("concat_and_cache_mla_cpu", torch::kCPU, &concat_and_cache_mla_cpu);
586 #endif
587
588 #if (defined(__AVX512BF16__) && defined(__AVX512F__) && \
tests/kernels/attention/test_amx_mla.py new
+471
@@ -0,0 +1,471 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Correctness tests for the AMX-only CPU MLA backend: the vendored
4 +decode/extend/bmm kernels, the KV cache write, and the ``AMXMLAImpl``
5 +backend built on top of them.
6 +"""
7 +
8 +import pytest
9 +import torch
10 +
11 +from vllm.platforms import current_platform
12 +
13 +if not current_platform.is_cpu():
14 + pytest.skip("skipping CPU-only tests", allow_module_level=True)
15 +
16 +if not torch.cpu._is_amx_tile_supported():
17 + pytest.skip("AMX MLA requires an AMX-capable host", allow_module_level=True)
18 +
19 +torch.cpu._init_amx()
20 +
21 +from vllm import _custom_ops as ops # noqa: E402
22 +from vllm.utils.torch_utils import set_random_seed # noqa: E402
23 +from vllm.v1.attention.backends.mla.amx_mla import ( # noqa: E402
24 + AMXMLAImpl,
25 + _compute_num_kv_splits,
26 + _expand_block_table,
27 +)
28 +
29 +KV_LORA_RANK = 64
30 +QK_ROPE_HEAD_DIM = 32
31 +HEAD_SIZE = KV_LORA_RANK + QK_ROPE_HEAD_DIM # 96, kv-cache row width
32 +# The extend/decode kernels transpose by 512-bit lanes: both the cache row
33 +# width and the (kv_lora_rank-wide) value width must be 32-element aligned.
34 +DTYPE = torch.bfloat16
35 +ATOL = 2e-2
36 +RTOL = 2e-2
37 +# Backend-level tests chain absorb -> attend -> de-absorb through bf16,
38 +# accumulating more rounding error than a single kernel op.
39 +_IMPL_ATOL = 1.5e-1
40 +_IMPL_RTOL = 1.5e-1
41 +
42 +
43 +def _flatten_cache(kv_cache: torch.Tensor) -> torch.Tensor:
44 + """(num_blocks, block_size, head_size) -> (num_slots, head_size)."""
45 + return kv_cache.view(-1, kv_cache.size(-1))
46 +
47 +
48 +def _random_paged_cache(num_slots: int, head_size: int) -> torch.Tensor:
49 + return torch.randn(num_slots, head_size, dtype=DTYPE)
50 +
51 +
52 +def _ref_latent_attn(
53 + q: torch.Tensor,
54 + keys: torch.Tensor,
55 + scale: float,
56 + kv_lora_rank: int,
57 +) -> torch.Tensor:
58 + """Plain-PyTorch causal-free latent-space MQA attention for one request.
59 +
60 + q: (num_heads, head_size); keys: (ctx_len, head_size). Values are the
61 + first ``kv_lora_rank`` columns of the same latent rows (MLA's K/V
62 + aliasing), matching what decode_attention_cpu/extend_attention_cpu
63 + compute internally for the MLA-shaped case.
64 + """
65 + values = keys[:, :kv_lora_rank]
66 + scores = (q.float() @ keys.float().T) * scale
67 + probs = torch.softmax(scores, dim=-1)
68 + return (probs @ values.float()).to(q.dtype)
69 +
70 +
71 +def test_bmm_cpu_matches_torch_bmm():
72 + set_random_seed(0)
73 + # bmm_cpu requires the output's last dim (mat2's out-features) to be a
74 + # multiple of 32 (tinygemm's tile width); the contraction dim is free.
75 + n, b, p, l = 8, 32, 32, 32 # noqa: E741
76 + mat1 = torch.randn(n, b, p, dtype=DTYPE)
77 + mat2 = torch.randn(n, l, p, dtype=DTYPE) # (N, OUT, IN) Linear convention
78 +
79 + ref = torch.bmm(mat1.float(), mat2.float().transpose(1, 2)).to(DTYPE)
80 +
81 + out = torch.empty(n, b, l, dtype=DTYPE)
82 + ops.bmm_cpu(out, mat1, mat2, False, None)
83 + torch.testing.assert_close(out, ref, atol=ATOL, rtol=RTOL)
84 +
85 + packed = torch.ops._C.convert_weight_packed(mat2.contiguous())
86 + out_vnni = torch.empty(n, b, l, dtype=DTYPE)
87 + ops.bmm_cpu(out_vnni, mat1, packed, True, None)
88 + torch.testing.assert_close(out_vnni, ref, atol=ATOL, rtol=RTOL)
89 +
90 +
91 +def test_concat_and_cache_mla_round_trip():
92 + set_random_seed(0)
93 + num_tokens = 37
94 + num_blocks, block_size = 8, 16
95 + kv_cache = torch.zeros(num_blocks, block_size, HEAD_SIZE, dtype=DTYPE)
96 +
97 + kv_c_normed = torch.randn(num_tokens, KV_LORA_RANK, dtype=DTYPE)
98 + k_pe = torch.randn(num_tokens, QK_ROPE_HEAD_DIM, dtype=DTYPE)
99 + slot_mapping = torch.randperm(num_blocks * block_size)[:num_tokens].to(torch.int64)
100 +
101 + ops.amx_mla_concat_and_cache(kv_c_normed, k_pe, kv_cache, slot_mapping)
102 +
103 + flat = _flatten_cache(kv_cache)
104 + read_back = flat[slot_mapping]
105 + expected = torch.cat([kv_c_normed, k_pe], dim=-1)
106 + torch.testing.assert_close(read_back, expected)
107 +
108 + untouched = torch.ones(num_blocks * block_size, dtype=torch.bool)
109 + untouched[slot_mapping] = False
110 + assert (flat[untouched] == 0).all()
111 +
112 +
113 +@pytest.mark.parametrize("seq_lens", [[5], [1, 300, 33], [513, 1, 129, 7]])
114 +def test_decode_attention_cpu_matches_reference(seq_lens):
115 + set_random_seed(0)
116 + num_seqs = len(seq_lens)
117 + num_heads = 4
118 + block_size = 32
119 + max_len = max(seq_lens)
120 + max_blocks = (max_len + block_size - 1) // block_size
121 + num_blocks = num_seqs * max_blocks
122 +
123 + kv_cache = _random_paged_cache(num_blocks * block_size, HEAD_SIZE)
124 + block_table = torch.randperm(num_blocks).view(num_seqs, max_blocks).to(torch.int64)
125 + req_to_token = _expand_block_table(block_table, block_size)
126 + req_pool_indices = torch.arange(num_seqs, dtype=torch.int64)
127 + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int64)
128 + scale = HEAD_SIZE**-0.5
129 +
130 + q = torch.randn(num_seqs, num_heads, HEAD_SIZE, dtype=DTYPE)
131 + kv_cache_flat = kv_cache.view(-1, 1, HEAD_SIZE)
132 + v_buffer = kv_cache_flat[..., :KV_LORA_RANK]
133 +
134 + o = torch.zeros(num_seqs, num_heads, KV_LORA_RANK, dtype=DTYPE)
135 + attn_logits = torch.zeros(
136 + num_seqs, num_heads, 4, KV_LORA_RANK + 1, dtype=torch.float32
137 + )
138 + ops.cpu_mla_decode(
139 + q,
140 + kv_cache_flat,
141 + v_buffer,
142 + o,
143 + None,
144 + None,
145 + None,
146 + attn_logits,
147 + req_to_token,
148 + req_pool_indices,
149 + seq_lens_t,
150 + scale,
151 + 0.0,
152 + False,
153 + 0,
154 + None,
155 + None,
156 + )
157 +
158 + flat_cache = _flatten_cache(kv_cache)
159 + for i, seq_len in enumerate(seq_lens):
160 + token_ids = req_to_token[i, :seq_len]
161 + keys = flat_cache[token_ids]
162 + ref = _ref_latent_attn(q[i], keys, scale, KV_LORA_RANK)
163 + torch.testing.assert_close(o[i], ref, atol=ATOL, rtol=RTOL)
164 +
165 +
166 +@pytest.mark.parametrize(
167 + "prefix_lens,extend_lens",
168 + [
169 + ([0], [9]), # fresh prefill, no cached prefix
170 + ([64], [17]), # continuation of a cached prefix
171 + ([0, 40, 128], [23, 5, 31]), # mixed batch
172 + ],
173 +)
174 +def test_extend_attention_cpu_matches_reference(prefix_lens, extend_lens):
175 + set_random_seed(0)
176 + num_seqs = len(prefix_lens)
177 + num_heads = 4
178 + block_size = 32
179 + seq_lens = [p + e for p, e in zip(prefix_lens, extend_lens)]
180 + max_len = max(seq_lens)
181 + max_blocks = (max_len + block_size - 1) // block_size
182 + num_blocks = num_seqs * max_blocks
183 +
184 + kv_cache = _random_paged_cache(num_blocks * block_size, HEAD_SIZE)
185 + block_table = torch.randperm(num_blocks).view(num_seqs, max_blocks).to(torch.int64)
186 + req_to_token = _expand_block_table(block_table, block_size)
187 + req_pool_indices = torch.arange(num_seqs, dtype=torch.int64)
188 + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int64)
189 + extend_seq_lens_t = torch.tensor(extend_lens, dtype=torch.int64)
190 + extend_start_loc = torch.cumsum(
191 + torch.tensor([0, *extend_lens[:-1]], dtype=torch.int64), dim=0
192 + )
193 + scale = HEAD_SIZE**-0.5
194 + total_new_tokens = sum(extend_lens)
195 +
196 + # The new tokens' latent K/V, already written into the cache (mirrors
197 + # do_kv_cache_update running before forward_mha in real usage).
198 + k_extend_flat = torch.randn(total_new_tokens, HEAD_SIZE, dtype=DTYPE)
199 + flat_cache = _flatten_cache(kv_cache)
200 + for i in range(num_seqs):
201 + prefix_len = prefix_lens[i]
202 + ext_len = extend_lens[i]
203 + new_token_ids = req_to_token[i, prefix_len : prefix_len + ext_len]
204 + start = extend_start_loc[i].item()
205 + flat_cache[new_token_ids] = k_extend_flat[start : start + ext_len]
206 +
207 + q_extend = torch.randn(total_new_tokens, num_heads, HEAD_SIZE, dtype=DTYPE)
208 + k_extend = k_extend_flat.unsqueeze(1)
209 + v_extend = k_extend[..., :KV_LORA_RANK]
210 + kv_cache_flat = kv_cache.view(-1, 1, HEAD_SIZE)
211 + v_buffer = kv_cache_flat[..., :KV_LORA_RANK]
212 +
213 + o_extend = torch.empty(total_new_tokens, num_heads, KV_LORA_RANK, dtype=DTYPE)
214 + ops.cpu_mla_extend(
215 + q_extend,
216 + k_extend,
217 + v_extend,
218 + o_extend,
219 + kv_cache_flat,
220 + v_buffer,
221 + req_to_token,
222 + req_pool_indices,
223 + seq_lens_t,
224 + extend_seq_lens_t,
225 + extend_start_loc,
226 + max(extend_lens),
227 + scale,
228 + 0.0,
229 + False,
230 + 0,
231 + None,
232 + None,
233 + None,
234 + )
235 +
236 + for i in range(num_seqs):
237 + prefix_len = prefix_lens[i]
238 + ext_len = extend_lens[i]
239 + start = extend_start_loc[i].item()
240 + for t in range(ext_len):
241 + visible = prefix_len + t + 1
242 + token_ids = req_to_token[i, :visible]
243 + keys = flat_cache[token_ids]
244 + ref = _ref_latent_attn(q_extend[start + t], keys, scale, KV_LORA_RANK)
245 + torch.testing.assert_close(o_extend[start + t], ref, atol=ATOL, rtol=RTOL)
246 +
247 +
248 +class _FakeLinear:
249 + """Minimal duck-typed stand-in for kv_b_proj: only .weight and
250 + .quant_method are read by get_and_maybe_dequant_weights()."""
251 +
252 + def __init__(self, weight: torch.Tensor):
253 + self.weight = weight
254 + self.quant_method = None
255 +
256 +
257 +def _make_amx_mla_impl(num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank):
258 + weight = torch.randn(
259 + num_heads * (qk_nope_head_dim + v_head_dim), kv_lora_rank, dtype=DTYPE
260 + )
261 + kv_b_proj = _FakeLinear(weight)
262 + impl = AMXMLAImpl(
263 + num_heads=num_heads,
264 + head_size=kv_lora_rank + QK_ROPE_HEAD_DIM,
265 + scale=(qk_nope_head_dim + QK_ROPE_HEAD_DIM) ** -0.5,
266 + num_kv_heads=1,
267 + alibi_slopes=None,
268 + sliding_window=None,
269 + kv_cache_dtype="auto",
270 + logits_soft_cap=None,
271 + attn_type="decoder",
272 + kv_sharing_target_layer_name=None,
273 + q_lora_rank=None,
274 + kv_lora_rank=kv_lora_rank,
275 + qk_nope_head_dim=qk_nope_head_dim,
276 + qk_rope_head_dim=QK_ROPE_HEAD_DIM,
277 + qk_head_dim=qk_nope_head_dim + QK_ROPE_HEAD_DIM,
278 + v_head_dim=v_head_dim,
279 + kv_b_proj=kv_b_proj,
280 + )
281 + impl.process_weights_after_loading(DTYPE)
282 + return impl, weight
283 +
284 +
285 +class _FakePrefillMetadata:
286 + def __init__(
287 + self, block_table, cpu_seq_lens, query_start_loc, req_to_token, req_pool_indices
288 + ):
289 + self.block_table = block_table
290 + self.cpu_seq_lens = cpu_seq_lens
291 + query_start_loc_i64 = query_start_loc.to(torch.int64)
292 + extend_seq_lens = query_start_loc_i64[1:] - query_start_loc_i64[:-1]
293 + self.extend_seq_lens = extend_seq_lens
294 + self.extend_start_loc = query_start_loc_i64[:-1]
295 + self.max_len_extend = int(extend_seq_lens.max().item())
296 + self.req_to_token = req_to_token
297 + self.req_pool_indices = req_pool_indices
298 +
299 +
300 +class _FakeDecodeMetadata:
301 + def __init__(self, block_table, seq_lens, req_to_token, req_pool_indices):
302 + self.block_table = block_table
303 + self.seq_lens_i64 = seq_lens.to(torch.int64)
304 + self.req_to_token = req_to_token
305 + self.req_pool_indices = req_pool_indices
306 + self.num_kv_splits = _compute_num_kv_splits(
307 + int(seq_lens.max().item()), current_platform.num_compute_units()
308 + )
309 +
310 +
311 +class _FakeAttnMetadata:
312 + def __init__(self, decode=None, prefill=None, max_seq_len=0):
313 + self.decode = decode
314 + self.prefill = prefill
315 + self.max_seq_len = max_seq_len
316 +
317 +
318 +def test_amx_mla_impl_forward_mqa_matches_reference(default_vllm_config):
319 + """forward_mqa receives already-absorbed Q (as MLAAttention.forward_impl
320 + would produce via layer.W_UK_T) and must reproduce plain per-head MLA
321 + decode attention through W_UK/W_UV."""
322 + set_random_seed(1)
323 + num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank = 4, 32, 32, 64
324 + impl, kv_b_weight = _make_amx_mla_impl(
325 + num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank
326 + )
327 + w_uk, w_uv = kv_b_weight.T.view(
328 + kv_lora_rank, num_heads, qk_nope_head_dim + v_head_dim
329 + ).split([qk_nope_head_dim, v_head_dim], dim=-1)
330 +
331 + seq_lens = [1, 40, 5]
332 + num_seqs = len(seq_lens)
333 + block_size = 32
334 + max_blocks = (max(seq_lens) + block_size - 1) // block_size
335 + num_blocks = num_seqs * max_blocks
336 + head_size = kv_lora_rank + QK_ROPE_HEAD_DIM
337 +
338 + kv_cache = torch.randn(num_blocks, block_size, head_size, dtype=DTYPE)
339 + block_table = torch.randperm(num_blocks).view(num_seqs, max_blocks).to(torch.int64)
340 + req_to_token = _expand_block_table(block_table, block_size)
341 +
342 + q_nope = torch.randn(num_seqs, num_heads, qk_nope_head_dim, dtype=DTYPE)
343 + q_pe = torch.randn(num_seqs, num_heads, QK_ROPE_HEAD_DIM, dtype=DTYPE)
344 + # Pre-absorb, mirroring MLAAttention.forward_impl's own bmm against W_UK_T.
345 + ql_nope = torch.einsum("bnp,lnp->bnl", q_nope.float(), w_uk.float()).to(DTYPE)
346 +
347 + attn_metadata = _FakeAttnMetadata(
348 + decode=_FakeDecodeMetadata(
349 + block_table=block_table,
350 + seq_lens=torch.tensor(seq_lens, dtype=torch.int64),
351 + req_to_token=req_to_token,
352 + req_pool_indices=torch.arange(num_seqs, dtype=torch.int64),
353 + ),
354 + max_seq_len=max(seq_lens),
355 + )
356 + o, lse = impl.forward_mqa((ql_nope, q_pe), kv_cache, attn_metadata, layer=None)
357 + assert lse is None
358 + # De-absorb, mirroring MLAAttention.forward_impl's own bmm against
359 + # layer._v_up_proj, to get back into real per-head V space for
360 + # comparison against a from-scratch reference.
361 + out_real = torch.einsum("bnl,lnv->bnv", o.float(), w_uv.float())
362 +
363 + flat_cache = _flatten_cache(kv_cache)
364 + scale = impl.scale
365 + for i, seq_len in enumerate(seq_lens):
366 + token_ids = req_to_token[i, :seq_len]
367 + keys_latent = flat_cache[token_ids] # (L, head_size)
368 + k_nope_latent, k_pe = keys_latent.split(
369 + [kv_lora_rank, QK_ROPE_HEAD_DIM], dim=-1
370 + )
371 + k_real = torch.einsum("lp,pnd->nld", k_nope_latent.float(), w_uk.float())
372 + k_pe_b = k_pe.float().unsqueeze(0).expand(num_heads, -1, -1)
373 + k_full = torch.cat([k_real, k_pe_b], dim=-1) # (num_heads, L, qk_head_dim)
374 + q_full = torch.cat(
375 + [q_nope[i], q_pe[i]], dim=-1
376 + ).float() # (num_heads, qk_head_dim)
377 + scores = torch.einsum("nd,nld->nl", q_full, k_full) * scale
378 + probs = torch.softmax(scores, dim=-1)
379 + v_real = torch.einsum("lp,pnv->nlv", k_nope_latent.float(), w_uv.float())
380 + ref = torch.einsum("nl,nlv->nv", probs, v_real)
381 + torch.testing.assert_close(out_real[i], ref, atol=_IMPL_ATOL, rtol=_IMPL_RTOL)
382 +
383 +
384 +def test_amx_mla_impl_forward_mha_matches_reference(default_vllm_config):
385 + """forward_mha receives raw unabsorbed Q and must attend correctly for a
386 + prefill batch with a mix of fresh and continued (cached-prefix)
387 + sequences."""
388 + set_random_seed(2)
389 + num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank = 4, 32, 32, 64
390 + impl, kv_b_weight = _make_amx_mla_impl(
391 + num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank
392 + )
393 + w_uk, w_uv = kv_b_weight.T.view(
394 + kv_lora_rank, num_heads, qk_nope_head_dim + v_head_dim
395 + ).split([qk_nope_head_dim, v_head_dim], dim=-1)
396 +
397 + prefix_lens = [0, 20]
398 + extend_lens = [13, 7]
399 + seq_lens = [p + e for p, e in zip(prefix_lens, extend_lens)]
400 + num_seqs = len(prefix_lens)
401 + block_size = 32
402 + max_blocks = (max(seq_lens) + block_size - 1) // block_size
403 + num_blocks = num_seqs * max_blocks
404 + head_size = kv_lora_rank + QK_ROPE_HEAD_DIM
405 +
406 + kv_cache = torch.randn(num_blocks, block_size, head_size, dtype=DTYPE)
407 + block_table = torch.randperm(num_blocks).view(num_seqs, max_blocks).to(torch.int64)
408 + req_to_token = _expand_block_table(block_table, block_size)
409 + flat_cache = _flatten_cache(kv_cache)
410 +
411 + total_new_tokens = sum(extend_lens)
412 + query_start_loc = torch.tensor(
413 + [0, *torch.cumsum(torch.tensor(extend_lens), dim=0).tolist()],
414 + dtype=torch.int64,
415 + )
416 +
417 + q_nope = torch.randn(total_new_tokens, num_heads, qk_nope_head_dim, dtype=DTYPE)
418 + q_pe = torch.randn(total_new_tokens, num_heads, QK_ROPE_HEAD_DIM, dtype=DTYPE)
419 + q = torch.cat([q_nope, q_pe], dim=-1)
420 +
421 + kv_c_normed = torch.randn(total_new_tokens, kv_lora_rank, dtype=DTYPE)
422 + k_pe = torch.randn(total_new_tokens, 1, QK_ROPE_HEAD_DIM, dtype=DTYPE)
423 +
424 + # Write the new tokens' latent K/V into the cache ahead of time, mirroring
425 + # do_kv_cache_update running before forward_mha in real usage.
426 + for i in range(num_seqs):
427 + prefix_len = prefix_lens[i]
428 + ext_len = extend_lens[i]
429 + new_token_ids = req_to_token[i, prefix_len : prefix_len + ext_len]
430 + start = query_start_loc[i].item()
431 + flat_cache[new_token_ids, :kv_lora_rank] = kv_c_normed[start : start + ext_len]
432 + flat_cache[new_token_ids, kv_lora_rank:] = k_pe[start : start + ext_len, 0]
433 +
434 + attn_metadata = _FakeAttnMetadata(
435 + prefill=_FakePrefillMetadata(
436 + block_table=block_table,
437 + cpu_seq_lens=torch.tensor(seq_lens, dtype=torch.int64),
438 + query_start_loc=query_start_loc,
439 + req_to_token=req_to_token,
440 + req_pool_indices=torch.arange(num_seqs, dtype=torch.int64),
441 + )
442 + )
443 + output = torch.empty(total_new_tokens, num_heads * v_head_dim, dtype=DTYPE)
444 + impl.forward_mha(
445 + q, kv_c_normed, k_pe, kv_cache, attn_metadata, k_scale=None, output=output
446 + )
447 + out_view = output.view(total_new_tokens, num_heads, v_head_dim).float()
448 +
449 + scale = impl.scale
450 + for i in range(num_seqs):
451 + prefix_len = prefix_lens[i]
452 + ext_len = extend_lens[i]
453 + start = query_start_loc[i].item()
454 + for t in range(ext_len):
455 + visible = prefix_len + t + 1
456 + token_ids = req_to_token[i, :visible]
457 + keys_latent = flat_cache[token_ids]
458 + k_nope_latent, k_pe_ref = keys_latent.split(
459 + [kv_lora_rank, QK_ROPE_HEAD_DIM], dim=-1
460 + )
461 + k_real = torch.einsum("lp,pnd->nld", k_nope_latent.float(), w_uk.float())
462 + k_pe_b = k_pe_ref.float().unsqueeze(0).expand(num_heads, -1, -1)
463 + k_full = torch.cat([k_real, k_pe_b], dim=-1)
464 + q_full = torch.cat([q_nope[start + t], q_pe[start + t]], dim=-1).float()
465 + scores = torch.einsum("nd,nld->nl", q_full, k_full) * scale
466 + probs = torch.softmax(scores, dim=-1)
467 + v_real = torch.einsum("lp,pnv->nlv", k_nope_latent.float(), w_uv.float())
468 + ref = torch.einsum("nl,nlv->nv", probs, v_real)
469 + torch.testing.assert_close(
470 + out_view[start + t], ref, atol=_IMPL_ATOL, rtol=_IMPL_RTOL
471 + )
tests/v1/attention/test_mla_backends.py
+2
@@ -172,9 +172,11 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch):
172 layer.kv_b_proj.quant_method = None
173 layer.is_aiter_triton_fp4_bmm_enabled = False
174 layer.is_aiter_triton_fp8_bmm_enabled = False
175 + layer.is_amx_bmm_enabled = False
176 layer.dcp_q_replicate = False
177 layer.quant_config = None
178 layer.layer_name = "test"
179 + layer.impl = SimpleNamespace(process_weights_after_loading=lambda act_dtype: None)
180
181 monkeypatch.setattr(
182 mla_attention_module, "set_default_quant_scales", lambda *_, **__: None
tests/v1/attention/test_mla_prefill_selector.py
+1
@@ -63,6 +63,7 @@ class TestGetMLAPrefillBackend:
63
64 with patch("vllm.platforms.current_platform") as mock_platform:
65 mock_platform.get_device_capability.return_value = None
66 + mock_platform.is_cpu.return_value = False
67
68 backend = get_mla_prefill_backend(vllm_config)
69 assert backend.get_name() == "FLASH_ATTN"
vllm/_custom_ops.py
+103
@@ -3937,6 +3937,109 @@ def cpu_attention_with_kv_cache(
3937 )
3938
3939
3940 +def cpu_mla_decode(
3941 + query: torch.Tensor,
3942 + k_buffer: torch.Tensor,
3943 + v_buffer: torch.Tensor,
3944 + output: torch.Tensor,
3945 + key: torch.Tensor | None,
3946 + value: torch.Tensor | None,
3947 + loc: torch.Tensor | None,
3948 + attn_logits: torch.Tensor,
3949 + req_to_token: torch.Tensor,
3950 + req_pool_indices: torch.Tensor,
3951 + seq_lens: torch.Tensor,
3952 + sm_scale: float,
3953 + logit_cap: float,
3954 + is_cross_attn: bool,
3955 + sliding_window_size: int,
3956 + encoder_lens: torch.Tensor | None,
3957 + sinks: torch.Tensor | None,
3958 +) -> None:
3959 + torch.ops._C.decode_attention_cpu(
3960 + query,
3961 + k_buffer,
3962 + v_buffer,
3963 + output,
3964 + key,
3965 + value,
3966 + loc,
3967 + attn_logits,
3968 + req_to_token,
3969 + req_pool_indices,
3970 + seq_lens,
3971 + sm_scale,
3972 + logit_cap,
3973 + is_cross_attn,
3974 + sliding_window_size,
3975 + encoder_lens,
3976 + sinks,
3977 + )
3978 +
3979 +
3980 +def cpu_mla_extend(
3981 + q_extend: torch.Tensor,
3982 + k_extend: torch.Tensor | None,
3983 + v_extend: torch.Tensor | None,
3984 + o_extend: torch.Tensor,
3985 + k_buffer: torch.Tensor,
3986 + v_buffer: torch.Tensor,
3987 + req_to_token: torch.Tensor,
3988 + req_pool_indices: torch.Tensor,
3989 + seq_lens: torch.Tensor,
3990 + extend_seq_lens: torch.Tensor,
3991 + extend_start_loc: torch.Tensor,
3992 + max_len_extend: int,
3993 + sm_scale: float,
3994 + logit_cap: float,
3995 + is_cross_attn: bool,
3996 + sliding_window_size: int,
3997 + encoder_lens: torch.Tensor | None,
3998 + sinks: torch.Tensor | None,
3999 + tree_mask: torch.Tensor | None = None,
4000 +) -> None:
4001 + torch.ops._C.extend_attention_cpu(
4002 + q_extend,
4003 + k_extend,
4004 + v_extend,
4005 + o_extend,
4006 + k_buffer,
4007 + v_buffer,
4008 + req_to_token,
4009 + req_pool_indices,
4010 + seq_lens,
4011 + extend_seq_lens,
4012 + extend_start_loc,
4013 + max_len_extend,
4014 + sm_scale,
4015 + logit_cap,
4016 + is_cross_attn,
4017 + sliding_window_size,
4018 + encoder_lens,
4019 + sinks,
4020 + tree_mask,
4021 + )
4022 +
4023 +
4024 +def bmm_cpu(
4025 + out: torch.Tensor,
4026 + mat1: torch.Tensor,
4027 + mat2: torch.Tensor,
4028 + is_vnni: bool,
4029 + scale: torch.Tensor | None = None,
4030 +) -> None:
4031 + torch.ops._C.bmm_cpu(out, mat1, mat2, is_vnni, scale)
4032 +
4033 +
4034 +def amx_mla_concat_and_cache(
4035 + kv_c_normed: torch.Tensor,
4036 + k_pe: torch.Tensor,
4037 + kv_cache: torch.Tensor,
4038 + slot_mapping: torch.Tensor,
4039 +) -> None:
4040 + torch.ops._C.concat_and_cache_mla_cpu(kv_c_normed, k_pe, kv_cache, slot_mapping)
4041 +
4042 +
4043 def cpu_gemm_wna16(
4044 input: torch.Tensor,
4045 q_weight: torch.Tensor,
vllm/model_executor/layers/attention/mla_attention.py
+41
@@ -540,6 +540,11 @@ class MLAAttention(nn.Module, AttentionLayerBase):
540 **extra_impl_args,
541 )
542 self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None)
543 + self.is_amx_bmm_enabled = getattr(self.impl, "uses_amx_bmm", False)
544 + # AMX reads kv_b_proj's weight directly and never calls it live; the
545 + # reference CPU MLA backend calls it but isn't perf-critical. Skip
546 + # the packed-kernel dispatch either way.
547 + kv_b_proj._cpu_skip_gemm_dispatch = True
548 self.use_direct_call = not current_platform.opaque_attention_op()
549
550 vllm_config = get_current_vllm_config()
@@ -908,6 +913,20 @@ class MLAAttention(nn.Module, AttentionLayerBase):
913 group_size=128,
914 transpose_bm=True,
915 )
916 + elif self.is_amx_bmm_enabled:
917 + # bmm_cpu computes out[n] = mat1[n] @ mat2[n]^T against
918 + # AMXMLAImpl's own (N, L, P) packed W_UK -- same as prefill.
919 + N, B, P = mqa_q_nope.shape
920 + L = self.kv_lora_rank
921 + mqa_ql_nope = mqa_q_nope.new_empty((N, B, L))
922 + ops.bmm_cpu(
923 + mqa_ql_nope,
924 + mqa_q_nope,
925 + self.impl._w_uk_packed, # type: ignore[attr-defined]
926 + True,
927 + None,
928 + )
929 + mqa_ql_nope = mqa_ql_nope.transpose(0, 1)
930 else:
931 # Pads the head_dim if necessary (for the underlying kernel)
932 N, B, P = mqa_q_nope.shape
@@ -1026,6 +1045,18 @@ class MLAAttention(nn.Module, AttentionLayerBase):
1045 return output_padded
1046
1047 def process_weights_after_loading(self, act_dtype: torch.dtype):
1048 + # Let per-backend impls do their own weight packing first (no-op
1049 + # unless overridden), mirroring Attention.process_weights_after_loading.
1050 + self.impl.process_weights_after_loading(act_dtype)
1051 +
1052 + if self.is_amx_bmm_enabled:
1053 + # AMXMLAImpl already packed its own W_UK/W_UV above, for both
1054 + # prefill and decode. Release the now-unused raw weight.
1055 + self.kv_b_proj.weight = torch.nn.Parameter(
1056 + torch.empty(0), requires_grad=False
1057 + )
1058 + return
1059 +
1060 # we currently do not have quantized bmm's which are needed for
1061 # `W_UV` and `W_UK_T`, we just store fp16/bf16 copies and perform
1062 # the bmm's in 16-bit, the extra memory overhead of this is fairly low
@@ -1192,6 +1223,16 @@ class MLAAttention(nn.Module, AttentionLayerBase):
1223 x = rocm_aiter_ops.triton_fp8_bmm(
1224 x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True, YQ=out
1225 )
1226 + elif self.is_amx_bmm_enabled:
1227 + # bmm_cpu computes out[n] = mat1[n] @ mat2[n]^T against
1228 + # AMXMLAImpl's own (N, V, L) packed W_UV -- same as prefill.
1229 + ops.bmm_cpu(
1230 + out.transpose(0, 1),
1231 + x,
1232 + self.impl._w_uv_packed, # type: ignore[attr-defined]
1233 + True,
1234 + None,
1235 + )
1236 else:
1237 # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)
1238 torch.bmm(x, self.W_UV, out=out.transpose(0, 1))
vllm/model_executor/layers/linear.py
+6
@@ -197,6 +197,12 @@ class UnquantizedLinearMethod(LinearMethodBase):
197
198 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
199 if current_platform.is_cpu():
200 + # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): not
201 + # perf-critical, so skip packing and use a plain fallback.
202 + if getattr(layer, "_cpu_skip_gemm_dispatch", False):
203 + layer.cpu_linear = torch.nn.functional.linear
204 + return
205 +
206 from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm
207
208 dispatch_cpu_unquantized_gemm(layer, remove_weight=True)
vllm/platforms/cpu.py
+31 -5
@@ -89,13 +89,30 @@ class CpuPlatform(Platform):
89 if attn_selector_config.use_sparse:
90 raise NotImplementedError("Sparse Attention is not supported on CPU.")
91 if attn_selector_config.use_mla:
92 + amx_available = (
93 + cls.get_cpu_architecture() == CpuArchEnum.X86
94 + and torch.cpu._is_amx_tile_supported()
95 + )
96 + if amx_available and selected_backend != AttentionBackendEnum.CPU_MLA:
97 + # Prefer AMX when available, unless CPU_MLA was explicitly requested.
98 + if (
99 + selected_backend
100 + and selected_backend != AttentionBackendEnum.AMX_MLA
101 + ):
102 + logger.info("Cannot use %s backend on CPU.", selected_backend)
103 + logger.info_once("Using %s backend.", AttentionBackendEnum.AMX_MLA.name)
104 + return AttentionBackendEnum.AMX_MLA.get_path()
105 # Reference MLA implementation on CPU. Performance is not the
106 # goal here; the backend simply wires the CPU decode kernel
107 # (`mla_decode_kvcache`) and an SDPA-based prefill together with
108 # the shared MLA scaffolding so that DeepSeek-style models can
109 # execute on CPU.
97 - if selected_backend and selected_backend != AttentionBackendEnum.CPU_MLA:
110 + if selected_backend and selected_backend not in (
111 + AttentionBackendEnum.CPU_MLA,
112 + AttentionBackendEnum.AMX_MLA,
113 + ):
114 logger.info("Cannot use %s backend on CPU.", selected_backend)
115 + logger.info_once("Using %s backend.", AttentionBackendEnum.CPU_MLA.name)
116 return AttentionBackendEnum.CPU_MLA.get_path()
117 if selected_backend and selected_backend != AttentionBackendEnum.CPU_ATTN:
118 logger.info("Cannot use %s backend on CPU.", selected_backend)
@@ -134,11 +151,20 @@ class CpuPlatform(Platform):
151 # The CPU MLA decode kernel only compiles with block_size=16 today
152 # (see csrc/cpu/mla_decode.cpp). If the model uses MLA we override
153 # the default block size regardless of user preference to avoid a
137 - # runtime kernel dispatch failure.
154 + # runtime kernel dispatch failure. AMX MLA has no such constraint
155 + # (same AMX-available condition as get_attn_backend_cls), so it's
156 + # excluded from this override.
157 cpu_mla_enabled = model_config is not None and getattr(
158 model_config, "use_mla", False
159 )
141 - if cpu_mla_enabled:
160 + amx_mla_enabled = (
161 + cpu_mla_enabled
162 + and cls.get_cpu_architecture() == CpuArchEnum.X86
163 + and torch.cpu._is_amx_tile_supported()
164 + and vllm_config.attention_config.backend != AttentionBackendEnum.CPU_MLA
165 + )
166 + reference_cpu_mla_enabled = cpu_mla_enabled and not amx_mla_enabled
167 + if reference_cpu_mla_enabled:
168 if cache_config.user_specified_block_size and cache_config.block_size != 16:
169 logger.warning(
170 "CPU MLA backend requires block_size=16, overriding "
@@ -149,7 +175,7 @@ class CpuPlatform(Platform):
175 elif not cache_config.user_specified_block_size:
176 cache_config.block_size = 128
177
152 - if not cpu_mla_enabled and cache_config.block_size % 32 != 0:
178 + if not reference_cpu_mla_enabled and cache_config.block_size % 32 != 0:
179 logger.warning(
180 "CPU backend prefers block_size is multiples of 32, "
181 "otherwise the performance is not optimized."
@@ -342,7 +368,7 @@ class CpuPlatform(Platform):
368 vllm_config.parallel_config.tensor_parallel_size
369 )
370
345 - if model_config is not None and model_config.use_mla:
371 + if model_config is not None and model_config.use_mla and not amx_mla_enabled:
372 logger.info_once(
373 "MLA is enabled on a non-GPU platform; forcing chunked "
374 "prefill and prefix caching to be disabled."
vllm/v1/attention/backends/mla/amx_mla.py new
+405
@@ -0,0 +1,405 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""AMX-only, high-performance MLA backend for DeepSeek V2/V3/R1 on CPU.
4 +
5 +Built on the AMX decode/extend/bmm kernels vendored under
6 +``csrc/cpu/sgl-kernels/``, plugged into vLLM's ``MLACommonBackend``/
7 +``MLACommonImpl`` abstraction the same way every other concrete MLA backend
8 +(``TritonMLAImpl``, etc.) does.
9 +
10 +This is a separate backend from the reference ``CPUMLABackend``
11 +(``vllm/v1/attention/backends/mla/cpu_mla.py``): that one targets every CPU
12 +(any dtype, block_size=16, ``mla_decode_kvcache`` decode kernel + inherited
13 +SDPA-style prefill) as a functional/CI reference, not performance. This
14 +backend instead requires AMX (bf16 only, block_size a multiple of 32) and is
15 +selected by the platform layer in preference to the reference backend
16 +whenever the host supports it -- see ``CpuPlatform.get_attn_backend_cls``.
17 +
18 +Two points where this backend differs structurally from the GPU backends,
19 +both explained in the CPU MLA design plan:
20 +
21 +- ``forward_mha`` is fully overridden (not inherited from
22 + ``MLACommonBaseImpl``): the GPU "compute-friendly" prefill path depends on
23 + a pluggable ``MLAPrefillBackend`` and CUDA-only chunked-context gather ops,
24 + neither of which exist on CPU. Instead, this attends directly in
25 + latent-MQA space via the ``extend_attention_cpu`` kernel, which handles
26 + cached-prefix continuation and fresh prefill in one causal pass (the KV
27 + cache already contains the new tokens by the time this runs, since
28 + ``do_kv_cache_update`` executes first).
29 +- Weight absorption for the prefill path is done with this impl's own
30 + VNNI-packed copies of ``W_UK``/``W_UV`` (computed in
31 + ``process_weights_after_loading``, using ``bmm_cpu``), rather than reading
32 + the generic ``layer.W_UK_T``/``layer.W_UV``/``layer._v_up_proj`` the way
33 + GPU backends do, since ``forward_mha``'s abstract signature has no
34 + ``layer`` parameter (only ``forward_mqa`` does) and that signature is left
35 + untouched. The decode path needs no such packing: ``MLAAttention.forward_impl``
36 + already absorbs/de-absorbs Q around the ``forward_mqa`` call using the
37 + generic (unpacked) ``layer.W_UK_T``/``layer._v_up_proj``, so
38 + ``forward_mqa`` here only has to invoke the decode kernel.
39 +"""
40 +
41 +from __future__ import annotations
42 +
43 +from typing import ClassVar
44 +
45 +import torch
46 +
47 +from vllm import _custom_ops as ops
48 +from vllm.model_executor.layers.attention.mla_attention import (
49 + MLACommonBackend,
50 + MLACommonImpl,
51 + MLACommonMetadata,
52 + MLACommonMetadataBuilder,
53 +)
54 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
55 + get_and_maybe_dequant_weights,
56 +)
57 +from vllm.platforms import current_platform
58 +from vllm.v1.attention.backend import AttentionLayer, AttentionType, MultipleOf
59 +
60 +_MIN_WORK_PER_SPLIT = 512
61 +_SPLIT_OCCUPANCY_MULTIPLIER = 2
62 +
63 +
64 +def _compute_num_kv_splits(max_seq_len: int, num_threads: int) -> int:
65 + """Mirrors TritonMLAImpl's _compute_num_kv_splits, using the CPU thread
66 + count in place of SM count."""
67 + ideal_splits = 1
68 + while ideal_splits < max(1, max_seq_len // _MIN_WORK_PER_SPLIT):
69 + ideal_splits *= 2
70 + max_splits = num_threads * _SPLIT_OCCUPANCY_MULTIPLIER
71 + return min(ideal_splits, max_splits)
72 +
73 +
74 +def _expand_block_table(block_table: torch.Tensor, block_size: int) -> torch.Tensor:
75 + """Adapter: vLLM's block-table paging -> the flat per-(request, position)
76 + physical row index the decode/extend kernels expect. Called once per step
77 + from ``AMXMLAMetadataBuilder.build``, not per layer.
78 + """
79 + offsets = torch.arange(
80 + block_size, device=block_table.device, dtype=block_table.dtype
81 + )
82 + flat = block_table.unsqueeze(-1) * block_size + offsets
83 + return flat.reshape(block_table.size(0), -1).contiguous()
84 +
85 +
86 +class AMXMLABackend(MLACommonBackend):
87 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
88 +
89 + @classmethod
90 + def get_supported_head_sizes(cls) -> list[int]:
91 + # kv_lora_rank(512) + qk_rope_head_dim(64), DeepSeek V2/V3/R1/V3.2.
92 + return [576]
93 +
94 + @staticmethod
95 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
96 + return [MultipleOf(32)]
97 +
98 + @classmethod
99 + def supports_block_size(cls, block_size: int | None) -> bool:
100 + if block_size is None:
101 + return True
102 + return block_size % 32 == 0
103 +
104 + @staticmethod
105 + def get_name() -> str:
106 + return "AMX_MLA"
107 +
108 + @staticmethod
109 + def get_impl_cls() -> type[AMXMLAImpl]:
110 + return AMXMLAImpl
111 +
112 + @staticmethod
113 + def get_builder_cls() -> type[AMXMLAMetadataBuilder]:
114 + return AMXMLAMetadataBuilder
115 +
116 +
117 +class AMXMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
118 + def build(self, common_prefix_len, common_attn_metadata, fast_build: bool = False):
119 + attn_metadata = super().build(
120 + common_prefix_len, common_attn_metadata, fast_build
121 + )
122 + block_size = self.kv_cache_spec.block_size
123 + if attn_metadata.decode is not None:
124 + # Built once per step and reused by every layer, instead of
125 + # recomputing the same tensor arithmetic in every forward_mqa call.
126 + decode = attn_metadata.decode
127 + decode.req_to_token = _expand_block_table( # type: ignore[attr-defined]
128 + decode.block_table, block_size
129 + )
130 + decode.seq_lens_i64 = decode.seq_lens.to(torch.int64) # type: ignore[attr-defined]
131 + decode.req_pool_indices = torch.arange( # type: ignore[attr-defined]
132 + decode.block_table.size(0),
133 + dtype=torch.int64,
134 + device=decode.block_table.device,
135 + )
136 + decode.num_kv_splits = _compute_num_kv_splits( # type: ignore[attr-defined]
137 + attn_metadata.max_seq_len, current_platform.num_compute_units()
138 + )
139 + if attn_metadata.prefill is not None:
140 + # cpu_seq_lens: total per-request context length (prefix + new
141 + # tokens), the one thing the extend kernel needs that isn't
142 + # already on the shared MLACommonPrefillMetadata.
143 + prefill = attn_metadata.prefill
144 + num_decodes = attn_metadata.num_decodes
145 + num_prefills = attn_metadata.num_prefills
146 + prefill.cpu_seq_lens = common_attn_metadata.seq_lens[ # type: ignore[attr-defined]
147 + num_decodes : num_decodes + num_prefills
148 + ].to(torch.int64)
149 + prefill.req_to_token = _expand_block_table( # type: ignore[attr-defined]
150 + prefill.block_table, block_size
151 + ).to(torch.int64)
152 + prefill.req_pool_indices = torch.arange( # type: ignore[attr-defined]
153 + prefill.block_table.size(0),
154 + dtype=torch.int64,
155 + device=prefill.block_table.device,
156 + )
157 + query_start_loc_i64 = prefill.query_start_loc.to(torch.int64)
158 + extend_seq_lens = query_start_loc_i64[1:] - query_start_loc_i64[:-1]
159 + prefill.extend_seq_lens = extend_seq_lens # type: ignore[attr-defined]
160 + prefill.extend_start_loc = query_start_loc_i64[:-1] # type: ignore[attr-defined]
161 + prefill.max_len_extend = int(extend_seq_lens.max().item()) # type: ignore[attr-defined]
162 + return attn_metadata
163 +
164 +
165 +class AMXMLAImpl(MLACommonImpl[MLACommonMetadata]):
166 + # Tells MLAAttention (mla_attention.py) to use this impl's own packed
167 + # W_UK/W_UV + bmm_cpu for the shared decode absorb/de-absorb, mirroring
168 + # is_aiter_triton_fp8_bmm_enabled/is_aiter_triton_fp4_bmm_enabled.
169 + uses_amx_bmm = True
170 +
171 + def __init__(
172 + self,
173 + num_heads: int,
174 + head_size: int,
175 + scale: float,
176 + num_kv_heads: int,
177 + alibi_slopes: list[float] | None,
178 + sliding_window: int | None,
179 + kv_cache_dtype: str,
180 + logits_soft_cap: float | None,
181 + attn_type: str,
182 + kv_sharing_target_layer_name: str | None,
183 + **mla_args,
184 + ) -> None:
185 + super().__init__(
186 + num_heads,
187 + head_size,
188 + scale,
189 + num_kv_heads,
190 + alibi_slopes,
191 + sliding_window,
192 + kv_cache_dtype,
193 + logits_soft_cap,
194 + attn_type,
195 + kv_sharing_target_layer_name,
196 + **mla_args,
197 + )
198 + unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
199 + if any(unsupported_features):
200 + raise NotImplementedError(
201 + "AMXMLAImpl does not support one of the following: "
202 + "alibi_slopes, sliding_window, logits_soft_cap"
203 + )
204 + if attn_type != AttentionType.DECODER:
205 + raise NotImplementedError(
206 + "Encoder self-attention and encoder/decoder cross-attention "
207 + "are not implemented for AMXMLAImpl"
208 + )
209 + self._w_uk_packed: torch.Tensor | None = None
210 + self._w_uv_packed: torch.Tensor | None = None
211 +
212 + def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
213 + # forward_mha (unlike forward_mqa) receives no `layer` argument, so it
214 + # cannot read the generic layer.W_UK_T/W_UV/._v_up_proj that
215 + # MLAAttention.process_weights_after_loading already computes. This
216 + # impl derives its own copies from self.kv_b_proj (already available
217 + # via the constructor) and VNNI-packs them once, ahead of time, for
218 + # use with bmm_cpu in forward_mha.
219 + kv_b_proj_weight = get_and_maybe_dequant_weights(
220 + self.kv_b_proj, out_dtype=act_dtype
221 + ).T
222 + kv_b_proj_weight = kv_b_proj_weight.view(
223 + self.kv_lora_rank, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
224 + )
225 + w_uk, w_uv = kv_b_proj_weight.split(
226 + [self.qk_nope_head_dim, self.v_head_dim], dim=-1
227 + )
228 + # bmm_cpu computes out[b] = mat1[b] @ mat2[b]^T (Linear-weight
229 + # convention: mat2[b] is (OUT, IN)).
230 + # absorb: ql_nope(N,B,L) = q_nope(N,B,P) @ w_uk_for_bmm(N,L,P)^T
231 + # de-absorb: out(N,B,V) = attn_out(N,B,L) @ w_uv_for_bmm(N,V,L)^T
232 + w_uk_for_bmm = w_uk.permute(1, 0, 2).contiguous() # (N, L, P)
233 + w_uv_for_bmm = w_uv.permute(1, 2, 0).contiguous() # (N, V, L)
234 + self._w_uk_packed = torch.ops._C.convert_weight_packed(w_uk_for_bmm)
235 + self._w_uv_packed = torch.ops._C.convert_weight_packed(w_uv_for_bmm)
236 +
237 + def do_kv_cache_update(
238 + self,
239 + kv_c_normed: torch.Tensor,
240 + k_pe: torch.Tensor,
241 + kv_cache: torch.Tensor,
242 + slot_mapping: torch.Tensor,
243 + kv_cache_dtype: str,
244 + k_scale: torch.Tensor,
245 + ) -> None:
246 + # Overrides the default (CUDA-only concat_and_cache_mla).
247 + if kv_cache.numel() == 0:
248 + return
249 + assert kv_cache_dtype == "auto", (
250 + "AMXMLAImpl only supports an unquantized (bf16) KV cache; fp8 "
251 + "KV cache for CPU MLA is not yet implemented."
252 + )
253 + ops.amx_mla_concat_and_cache(
254 + kv_c_normed, k_pe.squeeze(1), kv_cache, slot_mapping.flatten()
255 + )
256 +
257 + def forward_mqa(
258 + self,
259 + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
260 + kv_c_and_k_pe_cache: torch.Tensor,
261 + attn_metadata: MLACommonMetadata,
262 + layer: AttentionLayer,
263 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
264 + # q arrives already absorbed (tensor or (ql_nope, q_pe) tuple) --
265 + # MLAAttention.forward_impl does the Q-absorption bmm generically
266 + # before calling forward_mqa, using layer.W_UK_T. De-absorption is
267 + # likewise done generically afterwards via layer._v_up_proj. Nothing
268 + # left for this impl to do beyond invoking the decode kernel.
269 + if isinstance(q, tuple):
270 + q = torch.cat(q, dim=-1)
271 + assert attn_metadata.decode is not None
272 +
273 + num_tokens = q.shape[0]
274 + kv_cache_flat = kv_c_and_k_pe_cache.view(-1, 1, self.head_size)
275 + v_buffer = kv_cache_flat[..., : self.kv_lora_rank]
276 +
277 + seq_lens = attn_metadata.decode.seq_lens_i64 # type: ignore[attr-defined]
278 + req_to_token = attn_metadata.decode.req_to_token # type: ignore[attr-defined]
279 + req_pool_indices = attn_metadata.decode.req_pool_indices # type: ignore[attr-defined]
280 +
281 + num_kv_splits = attn_metadata.decode.num_kv_splits # type: ignore[attr-defined]
282 + o = torch.zeros(
283 + num_tokens,
284 + self.num_heads,
285 + self.kv_lora_rank,
286 + dtype=q.dtype,
287 + device=q.device,
288 + )
289 + attn_logits = torch.zeros(
290 + num_tokens,
291 + self.num_heads,
292 + num_kv_splits,
293 + self.kv_lora_rank + 1,
294 + dtype=torch.float32,
295 + device=q.device,
296 + )
297 +
298 + ops.cpu_mla_decode(
299 + q,
300 + kv_cache_flat,
301 + v_buffer,
302 + o,
303 + None,
304 + None,
305 + None, # loc: only meaningful when key/value are given; we
306 + # already wrote the new K/V via do_kv_cache_update before this.
307 + attn_logits,
308 + req_to_token,
309 + req_pool_indices,
310 + seq_lens,
311 + self.scale,
312 + 0.0,
313 + False,
314 + 0,
315 + None,
316 + None,
317 + )
318 + return o, None
319 +
320 + def forward_mha(
321 + self,
322 + q: torch.Tensor,
323 + kv_c_normed: torch.Tensor,
324 + k_pe: torch.Tensor,
325 + kv_c_and_k_pe_cache: torch.Tensor,
326 + attn_metadata: MLACommonMetadata,
327 + k_scale: torch.Tensor,
328 + output: torch.Tensor,
329 + output_scale: torch.Tensor | None = None,
330 + ) -> None:
331 + assert output_scale is None, (
332 + "AMXMLAImpl.forward_mha does not support fused output quantization"
333 + )
334 + prefill = attn_metadata.prefill
335 + assert prefill is not None
336 + assert self._w_uk_packed is not None and self._w_uv_packed is not None
337 +
338 + num_tokens = q.shape[0]
339 + num_heads = self.num_heads
340 +
341 + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
342 +
343 + # Absorb: (N, B, P) x (N, L, P)^T -> (N, B, L)
344 + # bmm_cpu tolerates non-contiguous mat1/out as long as their last dim
345 + # is contiguous, so q_nope's transpose and ql_nope's transposed view
346 + # can be passed directly -- no copy needed for either.
347 + q_nope_t = q_nope.transpose(0, 1)
348 + ql_nope = torch.empty(
349 + num_tokens, num_heads, self.kv_lora_rank, dtype=q.dtype, device=q.device
350 + )
351 + ops.bmm_cpu(ql_nope.transpose(0, 1), q_nope_t, self._w_uk_packed, True, None)
352 + mqa_q = torch.cat([ql_nope, q_pe], dim=-1)
353 +
354 + kv_cache_flat = kv_c_and_k_pe_cache.view(-1, 1, self.head_size)
355 + v_buffer = kv_cache_flat[..., : self.kv_lora_rank]
356 +
357 + req_to_token = prefill.req_to_token # type: ignore[attr-defined]
358 + req_pool_indices = prefill.req_pool_indices # type: ignore[attr-defined]
359 +
360 + seq_lens = prefill.cpu_seq_lens # type: ignore[attr-defined]
361 + extend_seq_lens = prefill.extend_seq_lens # type: ignore[attr-defined]
362 + extend_start_loc = prefill.extend_start_loc # type: ignore[attr-defined]
363 + max_len_extend = prefill.max_len_extend # type: ignore[attr-defined]
364 +
365 + # k_extend/v_extend: the new tokens' own latent K/V, aliased in one
366 + # 576-wide buffer (v_extend is a view of k_extend's first
367 + # kv_lora_rank columns) so the kernel can gather both from one read,
368 + # mirroring the cache's own layout.
369 + k_extend = torch.cat(
370 + [kv_c_normed, k_pe.reshape(num_tokens, self.qk_rope_head_dim)], dim=-1
371 + ).unsqueeze(1)
372 + v_extend = k_extend[..., : self.kv_lora_rank]
373 +
374 + attn_out = torch.empty(
375 + num_tokens, num_heads, self.kv_lora_rank, dtype=q.dtype, device=q.device
376 + )
377 + ops.cpu_mla_extend(
378 + mqa_q,
379 + k_extend,
380 + v_extend,
381 + attn_out,
382 + kv_cache_flat,
383 + v_buffer,
384 + req_to_token,
385 + req_pool_indices,
386 + seq_lens,
387 + extend_seq_lens,
388 + extend_start_loc,
389 + max_len_extend,
390 + self.scale,
391 + 0.0,
392 + False,
393 + 0,
394 + None,
395 + None,
396 + None,
397 + )
398 +
399 + # De-absorb directly into `output` (view, no extra copy on either side):
400 + # (N, B, L) x (N, V, L)^T -> (N, B, V)
401 + attn_out_t = attn_out.transpose(0, 1)
402 + output_view = output.view(num_tokens, num_heads, self.v_head_dim).transpose(
403 + 0, 1
404 + )
405 + ops.bmm_cpu(output_view, attn_out_t, self._w_uv_packed, True, None)
vllm/v1/attention/backends/mla/prefill/cpu_native.py new
+61
@@ -0,0 +1,61 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inert MLA prefill backend for the CPU platform.
4 +
5 +CPU's MLA attention impl (``CPUMLAImpl``) fully overrides ``forward_mha`` with a
6 +kernel that attends directly against the paged latent KV cache (covering both
7 +fresh prefill and cached-prefix continuation in one pass), so it never calls
8 +into a pluggable ``MLAPrefillBackend``'s ``run_prefill_new_tokens``/
9 +``run_prefill_context_chunk``. This class exists only to satisfy
10 +``MLACommonMetadataBuilder``'s structural requirement that
11 +``MLAAttention.prefill_backend`` be a real, cloneable object.
12 +"""
13 +
14 +from typing import TYPE_CHECKING
15 +
16 +import torch
17 +
18 +from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
19 +
20 +if TYPE_CHECKING:
21 + from vllm.model_executor.layers.attention.mla_attention import (
22 + MLACommonPrefillMetadata,
23 + )
24 +
25 +
26 +class CPUNativeMLAPrefillBackend(MLAPrefillBackend):
27 + """Placeholder MLA prefill backend for CPU; never actually invoked."""
28 +
29 + @staticmethod
30 + def get_name() -> str:
31 + return "CPU_NATIVE"
32 +
33 + def run_prefill_new_tokens(
34 + self,
35 + q: torch.Tensor,
36 + k: torch.Tensor,
37 + v: torch.Tensor,
38 + return_softmax_lse: bool,
39 + out: torch.Tensor | None = None,
40 + output_scale: torch.Tensor | None = None,
41 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
42 + raise AssertionError(
43 + "CPUNativeMLAPrefillBackend.run_prefill_new_tokens is unreachable: "
44 + "CPUMLAImpl.forward_mha fully overrides the dense-MHA prefill path "
45 + "and never calls the pluggable prefill backend."
46 + )
47 +
48 + def run_prefill_context_chunk(
49 + self,
50 + chunk: "MLACommonPrefillMetadata.ContextChunk",
51 + q: torch.Tensor,
52 + k: torch.Tensor,
53 + v: torch.Tensor,
54 + out: torch.Tensor | None = None,
55 + ) -> tuple[torch.Tensor, torch.Tensor]:
56 + raise AssertionError(
57 + "CPUNativeMLAPrefillBackend.run_prefill_context_chunk is "
58 + "unreachable: CPUMLAImpl.forward_mha fully overrides the "
59 + "dense-MHA prefill path and never calls the pluggable prefill "
60 + "backend."
61 + )
vllm/v1/attention/backends/mla/prefill/registry.py
+3
@@ -52,6 +52,9 @@ class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta):
52 "vllm.v1.attention.backends.mla.prefill.aiter_flash_attn."
53 "AiterFlashAttnPrefillBackend"
54 )
55 + CPU_NATIVE = (
56 + "vllm.v1.attention.backends.mla.prefill.cpu_native.CPUNativeMLAPrefillBackend"
57 + )
58 # Placeholder for third-party/custom backends - must be registered before use
59 # set to None to avoid alias with other backend, whose value is an empty string
60 CUSTOM = None
vllm/v1/attention/backends/mla/prefill/selector.py
+2
@@ -109,6 +109,8 @@ def get_mla_prefill_backend(
109
110 device_capability = current_platform.get_device_capability()
111 if device_capability is None:
112 + if current_platform.is_cpu():
113 + return MLAPrefillBackendEnum.CPU_NATIVE.get_class()
114 logger.info_once(
115 "Device capability not available, using FlashAttention MLA prefill backend."
116 )
vllm/v1/attention/backends/registry.py
+1
@@ -123,6 +123,7 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
123 )
124 CPU_ATTN = "vllm.v1.attention.backends.cpu_attn.CPUAttentionBackend"
125 CPU_MLA = "vllm.v1.attention.backends.mla.cpu_mla.CPUMLABackend"
126 + AMX_MLA = "vllm.v1.attention.backends.mla.amx_mla.AMXMLABackend"
127 TURBOQUANT = "vllm.v1.attention.backends.turboquant_attn.TurboQuantAttentionBackend"
128 # Placeholder for third-party/custom backends - must be registered before use
129 # set to None to avoid alias with other backend, whose value is an empty string